From 96cad1f8e21ca3cb6852374b0d5c349913e887b6 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 3 Sep 2026 08:52:12 +0100 Subject: [PATCH 1/4] fix(attachments): refuse unparseable chat attachments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A chat attachment with no parser fell through to SimpleDirectoryReader's plain-text open(), so a phone-uploaded video was "extracted" into megabytes of binary garbage, truncated, and stored with extraction.status == "ok". Gate attachments in two tiers instead. A suffix with a dedicated parser is admitted on its name — a PDF is binary and parses fine. Anything else has to read as text: the first 8KB are sampled and refused on a NUL byte or too many other control bytes. That keeps source, config and log files working through the plain-text fallthrough, and keeps out videos, archives and renamed binaries alike. The route checks the staged spool before anything is stored or queued; the worker repeats the check where the local file exists, raising the non-retryable AttachmentRejectedError. SUPPORTED_ATTACHMENT_EXTENSIONS gains the parser-backed suffixes it was missing (.tiff, .tif, .bmp, .webp, .vtt, .xml) and is now exactly the file extractor's keys plus .txt, with a test asserting the two agree. The composer applies the same rule client-side, so an unsupported file is named before it costs an upload, and a test pins the frontend list to the backend one. Attachment failures now show their reason inline under the chips rather than only in a hover tooltip, which a touch user can never see, and only after a send was attempted. Dropped `accept` from the dropzone: it discarded rejected drops with no feedback and disagreed with the server about text files. --- application/api/user/attachments/routes.py | 25 +++ application/parser/file/constants.py | 64 ++++++ application/upload_limits.py | 104 +++++++++ application/worker.py | 25 +++ frontend/src/components/MessageInput.tsx | 64 +++++- .../message-input/AttachFileButton.tsx | 13 +- .../message-input/AttachmentChipList.tsx | 200 ++++++++++-------- frontend/src/constants/fileUpload.test.ts | 157 ++++++++++++++ frontend/src/constants/fileUpload.ts | 126 +++++++++++ frontend/src/locale/en.json | 3 +- frontend/src/upload/uploadSlice.ts | 2 + tests/api/user/attachments/test_routes.py | 94 ++++++++ tests/parser/file/test_constants.py | 95 +++++++++ tests/test_upload_limits.py | 123 +++++++++++ tests/worker/test_attachment_worker.py | 83 ++++++++ 15 files changed, 1079 insertions(+), 99 deletions(-) create mode 100644 frontend/src/constants/fileUpload.test.ts create mode 100644 tests/parser/file/test_constants.py diff --git a/application/api/user/attachments/routes.py b/application/api/user/attachments/routes.py index f3e7a8b5..e8d7491d 100644 --- a/application/api/user/attachments/routes.py +++ b/application/api/user/attachments/routes.py @@ -39,6 +39,9 @@ from application.stt.stt_creator import STTCreator from application.tts.tts_creator import TTSCreator from application.upload_limits import ( copy_upload_to_path, + enforce_parseable_attachment, + is_unsupported_upload_message, + UnsupportedUploadTypeError, upload_limit_message, UploadTooLargeError, ) @@ -130,6 +133,8 @@ def _get_store_attachment_user_error(exc: Exception) -> str: return build_stt_file_size_limit_message() if isinstance(exc, UploadTooLargeError): return upload_limit_message() + if isinstance(exc, UnsupportedUploadTypeError): + return str(exc) return "Failed to process file" @@ -214,6 +219,10 @@ class StoreAttachment(Resource): with tempfile.TemporaryDirectory() as temp_dir: staged_path = os.path.join(temp_dir, original_filename) copy_upload_to_path(file, staged_path) + # Refuse here — nothing is stored or queued yet — so a + # video the picker let through never reaches the + # worker's reader, which would open it as text. + enforce_parseable_attachment(staged_path, original_filename) with open(staged_path, "rb") as staged_stream: staged_upload = FileStorage( stream=staged_stream, @@ -263,6 +272,22 @@ class StoreAttachment(Resource): ), 413, ) + if errors and all( + is_unsupported_upload_message(error.get("error")) for error in errors + ): + # Tell the user which type was refused rather than the + # generic copy — this is the branch a phone picker that + # ignores ``accept`` lands in. + return make_response( + jsonify( + { + "success": False, + "message": errors[0]["error"], + "errors": errors, + } + ), + 400, + ) return make_response( jsonify({"status": "error", "message": "No valid files to upload"}), 400, diff --git a/application/parser/file/constants.py b/application/parser/file/constants.py index 4656727b..079a13d2 100644 --- a/application/parser/file/constants.py +++ b/application/parser/file/constants.py @@ -1,5 +1,7 @@ """Shared file-extension constants for parsing and ingestion flows.""" +import os + from application.stt.constants import SUPPORTED_AUDIO_EXTENSIONS @@ -25,3 +27,65 @@ SUPPORTED_SOURCE_EXTENSIONS = ( *SUPPORTED_SOURCE_IMAGE_EXTENSIONS, *SUPPORTED_AUDIO_EXTENSIONS, ) + +# Suffixes the attachment path has a dedicated parser for — every key of +# ``get_default_file_extractor()`` plus ``.txt``, which needs none. Kept as a +# literal (importing ``bulk`` here would drag docling into the API process), +# with ``tests/parser/file/test_constants.py`` asserting the two agree. +# +# This is *not* the whole attachment allow-list: a suffix with no parser is +# read by ``SimpleDirectoryReader``'s plain-text fallthrough, which is right +# for a .py or a .log and catastrophic for a video — the reader "extracts" +# megabytes of binary garbage, truncates it, and stores it with +# ``extraction.status == "ok"``. So unparsed suffixes are admitted on +# content instead (``upload_limits.enforce_parseable_attachment``): text +# passes, binary is refused. Zip is deliberately absent — source ingestion +# extracts archives, the attachment path does not, and a zip fails the +# content check like any other binary. +# +# Mirrored in ``frontend/src/constants/fileUpload.ts``; update both together. +SUPPORTED_ATTACHMENT_EXTENSIONS = frozenset( + { + *SUPPORTED_SOURCE_EXTENSIONS, + ".xhtml", + ".adoc", + ".asciidoc", + ".tiff", + ".tif", + ".bmp", + ".webp", + ".vtt", + ".xml", + } +) + + +def attachment_extension(filename: str | None) -> str: + """Return the lower-cased extension of ``filename`` including the dot, or ``""``. + + Args: + filename: A bare filename or path; ``None`` and empty strings yield ``""``. + + Returns: + The last suffix in lower case (``".pdf"``), or ``""`` when there is none. + """ + if not filename: + return "" + return os.path.splitext(os.path.basename(str(filename)))[1].lower() + + +def has_attachment_parser(filename: str | None) -> bool: + """Return whether an attachment's suffix has a dedicated parser. + + A False result does not mean the file is refused: it means nothing but the + plain-text fallthrough will read it, so it has to earn its place on + content. The decision is by extension only, never by the mime type a + browser reports (mobile pickers ignore ``accept`` and lie). + + Args: + filename: The upload's filename. + + Returns: + True when the suffix is in ``SUPPORTED_ATTACHMENT_EXTENSIONS``. + """ + return attachment_extension(filename) in SUPPORTED_ATTACHMENT_EXTENSIONS diff --git a/application/upload_limits.py b/application/upload_limits.py index 98736c29..c5f5c579 100644 --- a/application/upload_limits.py +++ b/application/upload_limits.py @@ -8,6 +8,10 @@ from contextlib import suppress from typing import BinaryIO from application.core.settings import settings +from application.parser.file.constants import ( + attachment_extension, + has_attachment_parser, +) _COPY_CHUNK_BYTES = 64 * 1024 @@ -26,6 +30,106 @@ class UploadTooLargeError(ValueError): """Raised when one uploaded file exceeds the configured byte cap.""" +class UnsupportedUploadTypeError(ValueError): + """Raised when an uploaded attachment has a file type the worker cannot parse.""" + + +_UNSUPPORTED_UPLOAD_PREFIX = "Unsupported file type" + + +def unsupported_upload_message(filename: str | None) -> str: + """Return the stable client-facing rejection message for an unparseable upload. + + Args: + filename: The rejected upload's filename. + + Returns: + ``"Unsupported file type: .mp4"`` style text, naming the extension. + """ + extension = attachment_extension(filename) + return f"{_UNSUPPORTED_UPLOAD_PREFIX}: {extension or '(no extension)'}" + + +def is_unsupported_upload_message(message: str | None) -> bool: + """Return whether ``message`` is one produced by :func:`unsupported_upload_message`.""" + return bool(message) and str(message).startswith(_UNSUPPORTED_UPLOAD_PREFIX) + + +# A suffix with no parser is read by ``SimpleDirectoryReader``'s plain-text +# fallthrough, so it is admitted on content: enough of the head to recognise a +# container header, and a tolerance that keeps real text (UTF-8 accents, an +# ANSI-coloured log) in while a random binary — ~12.5% of bytes below 0x20 — +# stays out. +_TEXT_SNIFF_BYTES = 8192 +_MAX_NONTEXT_RATIO = 0.10 +# Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. +_TEXT_CONTROL_BYTES = frozenset({0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x1B}) + + +def looks_like_text(sample: bytes) -> bool: + """Return whether a leading byte sample reads as text rather than binary. + + Args: + sample: The first bytes of a file; an empty sample counts as text. + + Returns: + False when the sample holds a NUL byte or too many other non-text + control bytes, True otherwise. + """ + if not sample: + return True + if b"\x00" in sample: + return False + nontext = sum( + 1 + for byte in sample + if (byte < 0x20 and byte not in _TEXT_CONTROL_BYTES) or byte == 0x7F + ) + return nontext / len(sample) <= _MAX_NONTEXT_RATIO + + +def file_looks_like_text(path: str | os.PathLike[str]) -> bool: + """Return whether a file on disk reads as text, by its leading bytes. + + Args: + path: Filesystem path to sample. + + Returns: + The :func:`looks_like_text` verdict for the file's head; True when the + file cannot be read, leaving that failure to the parser to report. + """ + try: + with open(path, "rb") as handle: + return looks_like_text(handle.read(_TEXT_SNIFF_BYTES)) + except OSError: + return True + + +def enforce_parseable_attachment( + path: str | os.PathLike[str], filename: str | None +) -> None: + """Reject an attachment that no parser handles and that is not plain text. + + Suffixes with a parser (``SUPPORTED_ATTACHMENT_EXTENSIONS``) are admitted + unconditionally — a PDF is binary and parses fine. Everything else has to + read as text, which is what keeps a video or an archive out of the + plain-text fallthrough while leaving source, config and log files in. + + Args: + path: Local path to the staged upload, readable before it is stored. + filename: The upload's filename, used for the suffix and the message. + + Raises: + UnsupportedUploadTypeError: When the file has no parser and its + contents are binary. + """ + if has_attachment_parser(filename): + return + if file_looks_like_text(path): + return + raise UnsupportedUploadTypeError(unsupported_upload_message(filename)) + + class _LimitedRawReader(io.RawIOBase): """Expose a binary stream while enforcing an exact cumulative read cap.""" diff --git a/application/worker.py b/application/worker.py index 0e7551f0..a1507239 100755 --- a/application/worker.py +++ b/application/worker.py @@ -55,6 +55,10 @@ from application.storage.db.repositories.wiki_pages import ( from application.storage.db.session import db_readonly, db_session from application.storage.db.source_config import SourceConfig from application.storage.storage_creator import StorageCreator +from application.upload_limits import ( + enforce_parseable_attachment, + UnsupportedUploadTypeError, +) from application.utils import ( count_tokens_docs, get_encoding, @@ -1554,6 +1558,26 @@ class AttachmentRejectedError(Exception): """ +def _reject_unparseable_attachment(local_path: str, filename: str) -> None: + """Reject an attachment with no parser whose contents are binary. + + Defence in depth behind the route's gate: ``SimpleDirectoryReader`` opens + any suffix it has no parser for as plain text, so a binary must fail here, + visibly, rather than "succeed" as garbage. + + Args: + local_path: Filesystem path of the attachment about to be parsed. + filename: The upload's original filename, which carries the suffix. + + Raises: + AttachmentRejectedError: If the file has no parser and is not text. + """ + try: + enforce_parseable_attachment(local_path, filename) + except UnsupportedUploadTypeError as exc: + raise AttachmentRejectedError(str(exc)) from exc + + def _reject_attachment_zip_bomb(local_path: str) -> None: """Reject a zip-container attachment that decompresses to too much. @@ -1751,6 +1775,7 @@ def attachment_worker(self, file_info, user): parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader" def _parse_local_file(local_path: str, **kwargs) -> Document: + _reject_unparseable_attachment(local_path, filename) _reject_attachment_zip_bomb(local_path) parse_path, is_temp_copy = _bounded_attachment_copy(local_path) try: diff --git a/frontend/src/components/MessageInput.tsx b/frontend/src/components/MessageInput.tsx index 2dc6d364..fac83d66 100644 --- a/frontend/src/components/MessageInput.tsx +++ b/frontend/src/components/MessageInput.tsx @@ -48,7 +48,9 @@ import { useArmedSend } from './message-input/armedSend'; import { handleAbort } from '../conversation/conversationSlice'; import { AUDIO_FILE_ACCEPT_ATTR, - FILE_UPLOAD_ACCEPT, + getFileExtension, + parseUploadErrorMessage, + partitionAttachmentFiles, } from '../constants/fileUpload'; import { UserToolType } from '../settings/types'; import { isChatToolVisible } from '../utils/toolUtils'; @@ -497,8 +499,31 @@ export default function MessageInput({ ); const uploadFiles = useCallback( - (files: File[]) => { - if (!files || files.length === 0) return; + async (incomingFiles: File[]) => { + if (!incomingFiles || incomingFiles.length === 0) return; + + // Run the server's own rule here, not just the input's `accept`: + // mobile pickers ignore `accept`, and a file the server will refuse + // should say so before it costs an upload. Surface the refusal as a + // failed chip so the user sees why instead of a silent drop. + const { supported, unsupported } = + await partitionAttachmentFiles(incomingFiles); + unsupported.forEach((file) => { + dispatch( + addAttachment({ + id: generateId(), + fileName: file.name, + progress: 0, + status: 'failed' as const, + taskId: '', + errorMessage: t('conversation.attachments.unsupportedType', { + extension: getFileExtension(file.name) || '?', + }), + }), + ); + }); + if (supported.length === 0) return; + const files = supported; const apiHost = import.meta.env.VITE_API_HOST; @@ -579,9 +604,14 @@ export default function MessageInput({ tasksByIndex.set(uploadIndex, task); }); + const errorsByIndex = new Map(); errors.forEach((errorItem) => { if (typeof errorItem.upload_index === 'number') { failedIndices.add(errorItem.upload_index); + errorsByIndex.set( + errorItem.upload_index, + errorItem.error, + ); } }); @@ -616,7 +646,10 @@ export default function MessageInput({ dispatch( updateAttachment({ id: uiId, - updates: { status: 'failed' }, + updates: { + status: 'failed', + errorMessage: errorsByIndex.get(index), + }, }), ); return; @@ -748,11 +781,12 @@ export default function MessageInput({ } } else { console.error('Upload failed', status, xhr.responseText); + const errorMessage = parseUploadErrorMessage(xhr.responseText); Object.values(indexToUiId).forEach((id) => dispatch( updateAttachment({ id, - updates: { status: 'failed' }, + updates: { status: 'failed', errorMessage }, }), ), ); @@ -871,7 +905,10 @@ export default function MessageInput({ dispatch( updateAttachment({ id: uniqueId, - updates: { status: 'failed' }, + updates: { + status: 'failed', + errorMessage: parseUploadErrorMessage(xhr.responseText), + }, }), ); } @@ -891,7 +928,7 @@ export default function MessageInput({ xhr.send(formData); }); }, - [dispatch, token, trackAttachment], + [dispatch, t, token, trackAttachment], ); const handleFileAttachment = (e: React.ChangeEvent) => { @@ -928,7 +965,10 @@ export default function MessageInput({ setHandleDragActive(false); }, maxSize: 25000000, - accept: FILE_UPLOAD_ACCEPT, + // No `accept`: react-dropzone would drop a rejected file on the floor + // with no feedback, and its mime matching disagrees with the server for + // text files that have no parser (.py, .log). uploadFiles applies the + // server's rule and reports what it refuses. }); const handleInput = useCallback(() => { @@ -1594,7 +1634,13 @@ export default function MessageInput({ onChange={handleVoiceFileAttachment} /> -
+ {/* translate="no": keep Chrome's page translator out of the composer — + it rewrites text nodes into wrappers and React loses the + controls (dead Attach button, see AttachFileButton). */} +
) => void; @@ -11,7 +11,14 @@ export default function AttachFileButton({ onChange }: AttachFileButtonProps) { const { t } = useTranslation(); return ( -
+ ); + })} +
+ + {failures.length > 0 && ( +
+ {failures.map((attachment) => ( + + {attachment.fileName}: {attachment.errorMessage} + + ))} +
+ )} + ); } diff --git a/frontend/src/constants/fileUpload.test.ts b/frontend/src/constants/fileUpload.test.ts new file mode 100644 index 00000000..ff0b9769 --- /dev/null +++ b/frontend/src/constants/fileUpload.test.ts @@ -0,0 +1,157 @@ +import { describe, expect, it } from 'vitest'; + +import { + ATTACHMENT_FILE_ACCEPT_ATTR, + ATTACHMENT_PARSER_EXTENSIONS, + getFileExtension, + hasAttachmentParser, + looksLikeText, + parseUploadErrorMessage, + partitionAttachmentFiles, +} from './fileUpload'; + +const file = (name: string, body: BlobPart = 'x', type = '') => + new File([body], name, { type }); + +const bytes = (...values: number[]) => new Uint8Array(values); + +const MP4_HEADER = bytes( + 0x00, + 0x00, + 0x00, + 0x18, + 0x66, + 0x74, + 0x79, + 0x70, + 0x69, + 0x73, + 0x6f, + 0x6d, +); + +describe('attachment type gate', () => { + it('extracts a lower-cased extension', () => { + expect(getFileExtension('Photo.JPG')).toBe('.jpg'); + expect(getFileExtension('a.tar.gz')).toBe('.gz'); + expect(getFileExtension('noext')).toBe(''); + expect(getFileExtension('.env')).toBe(''); + }); + + it('recognises parser-backed suffixes by name, case-insensitively', () => { + expect(hasAttachmentParser(file('Report.PDF'))).toBe(true); + expect(hasAttachmentParser(file('scan.WebP'))).toBe(true); + expect(hasAttachmentParser(file('subs.vtt'))).toBe(true); + // No parser — these are judged on content, not on the suffix. + expect(hasAttachmentParser(file('main.py'))).toBe(false); + expect(hasAttachmentParser(file('archive.zip'))).toBe(false); + expect(hasAttachmentParser(file('clip.mp4'))).toBe(false); + }); + + it('mirrors the backend list: no zip, no video, images and markup included', () => { + const listed = ATTACHMENT_FILE_ACCEPT_ATTR.split(','); + expect(listed).toEqual([...ATTACHMENT_PARSER_EXTENSIONS]); + expect(listed).toContain('.pdf'); + // Parser-backed suffixes the first cut of this list missed. + for (const ext of [ + '.webp', + '.tiff', + '.tif', + '.bmp', + '.vtt', + '.xml', + '.mdx', + ]) { + expect(listed).toContain(ext); + } + expect(listed).not.toContain('.zip'); + expect(listed).not.toContain('.mp4'); + }); +}); + +describe('looksLikeText', () => { + it('accepts empty, plain, accented and ANSI-coloured text', () => { + expect(looksLikeText(new Uint8Array())).toBe(true); + expect(looksLikeText(new TextEncoder().encode('plain text\n'))).toBe(true); + expect(looksLikeText(new TextEncoder().encode('café — em dash\n'))).toBe( + true, + ); + expect(looksLikeText(bytes(0x1b, 0x5b, 0x33, 0x31, 0x6d, 0x6f, 0x6b))).toBe( + true, + ); + }); + + it('rejects a NUL byte and dense control bytes', () => { + expect(looksLikeText(MP4_HEADER)).toBe(false); + expect(looksLikeText(bytes(0x74, 0x78, 0x00, 0x74))).toBe(false); + expect( + looksLikeText( + new Uint8Array(Array.from({ length: 64 }, (_, i) => i + 1)), + ), + ).toBe(false); + }); +}); + +describe('partitionAttachmentFiles', () => { + it('refuses video and archives whatever the picker claims the type is', async () => { + const { supported, unsupported } = await partitionAttachmentFiles([ + file('clip.mp4', MP4_HEADER, 'text/plain'), + file('archive.zip', bytes(0x50, 0x4b, 0x03, 0x04, 0x00, 0x00)), + ]); + expect(supported).toEqual([]); + expect(unsupported.map((f) => f.name)).toEqual(['clip.mp4', 'archive.zip']); + }); + + it('keeps text files that have no parser — the backend reads them', async () => { + const { supported, unsupported } = await partitionAttachmentFiles([ + file('main.py', 'def main():\n return 1\n'), + file('server.log', '2026-09-02 ERROR boom\n'), + file('Dockerfile', 'FROM python:3.12\n'), + ]); + expect(supported.map((f) => f.name)).toEqual([ + 'main.py', + 'server.log', + 'Dockerfile', + ]); + expect(unsupported).toEqual([]); + }); + + it('admits parser-backed types on their suffix, binary contents and all', async () => { + // A PDF is binary; the sniff must never be applied to it. + const { supported, unsupported } = await partitionAttachmentFiles([ + file('paper.pdf', MP4_HEADER), + file('photo.jpeg', MP4_HEADER), + ]); + expect(supported.map((f) => f.name)).toEqual(['paper.pdf', 'photo.jpeg']); + expect(unsupported).toEqual([]); + }); + + it('partitions a mixed selection and keeps the original order', async () => { + const mp4 = file('clip.mp4', MP4_HEADER); + const txt = file('notes.txt'); + const pdf = file('paper.pdf'); + const { supported, unsupported } = await partitionAttachmentFiles([ + mp4, + txt, + pdf, + ]); + expect(supported).toEqual([txt, pdf]); + expect(unsupported).toEqual([mp4]); + }); +}); + +describe('parseUploadErrorMessage', () => { + it('returns the server message when present', () => { + expect( + parseUploadErrorMessage( + '{"success":false,"message":"Unsupported file type: .mp4"}', + ), + ).toBe('Unsupported file type: .mp4'); + }); + + it('returns undefined for non-JSON bodies or a missing message', () => { + expect(parseUploadErrorMessage('502')).toBeUndefined(); + expect(parseUploadErrorMessage('{"ok":1}')).toBeUndefined(); + expect(parseUploadErrorMessage('')).toBeUndefined(); + }); +}); diff --git a/frontend/src/constants/fileUpload.ts b/frontend/src/constants/fileUpload.ts index e8d6f501..b8fd7c0d 100644 --- a/frontend/src/constants/fileUpload.ts +++ b/frontend/src/constants/fileUpload.ts @@ -80,3 +80,129 @@ export const SOURCE_FILE_TREE_ACCEPT_ATTR = [ '.jpg', '.jpeg', ].join(','); + +/** + * Chat-attachment suffixes with a dedicated parser. Mirrors the backend's + * `SUPPORTED_ATTACHMENT_EXTENSIONS` (application/parser/file/constants.py) — + * update both together. Zip is absent: source ingestion extracts archives, + * the attachment path does not. + * + * Not the whole allow-list. A suffix that isn't here (.py, .log, .yaml) is + * read by the backend's plain-text fallthrough, so it is judged on content + * by `partitionAttachmentFiles` instead, exactly as the server does. + */ +export const ATTACHMENT_PARSER_EXTENSIONS: readonly string[] = [ + '.rst', + '.md', + '.mdx', + '.pdf', + '.txt', + '.docx', + '.csv', + '.epub', + '.html', + '.xhtml', + '.json', + '.xlsx', + '.pptx', + '.adoc', + '.asciidoc', + '.png', + '.jpg', + '.jpeg', + '.tiff', + '.tif', + '.bmp', + '.webp', + '.vtt', + '.xml', + '.wav', + '.mp3', + '.m4a', + '.ogg', + '.webm', +]; + +/** Picker filter for the Attach button. A hint only — pickers may ignore it. */ +export const ATTACHMENT_FILE_ACCEPT_ATTR = + ATTACHMENT_PARSER_EXTENSIONS.join(','); + +/** Lower-cased last extension including the dot, or '' (dotfiles have none). */ +export function getFileExtension(name: string): string { + const base = name.split(/[\\/]/).pop() ?? ''; + const dot = base.lastIndexOf('.'); + if (dot <= 0) return ''; + return base.slice(dot).toLowerCase(); +} + +/** + * Whether the backend has a parser for this suffix. False is not a refusal: + * the file then has to read as text. Never decided by the mime type the + * picker reports — mobile pickers ignore `accept` and lie about types. + */ +export function hasAttachmentParser(file: { name: string }): boolean { + const ext = getFileExtension(file.name); + return ext !== '' && ATTACHMENT_PARSER_EXTENSIONS.includes(ext); +} + +// Mirrors application/upload_limits.py: enough of the head to recognise a +// container header, and a tolerance that keeps real text (UTF-8 accents, an +// ANSI-coloured log) in while a random binary stays out. +const TEXT_SNIFF_BYTES = 8192; +const MAX_NONTEXT_RATIO = 0.1; +// Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. +const TEXT_CONTROL_BYTES = new Set([0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x1b]); + +/** Whether a leading byte sample reads as text rather than binary. */ +export function looksLikeText(sample: Uint8Array): boolean { + if (sample.length === 0) return true; + let nontext = 0; + for (const byte of sample) { + if (byte === 0x00) return false; + if ((byte < 0x20 && !TEXT_CONTROL_BYTES.has(byte)) || byte === 0x7f) + nontext += 1; + } + return nontext / sample.length <= MAX_NONTEXT_RATIO; +} + +async function isSupportedAttachmentFile(file: File): Promise { + if (hasAttachmentParser(file)) return true; + try { + const head = await file.slice(0, TEXT_SNIFF_BYTES).arrayBuffer(); + return looksLikeText(new Uint8Array(head)); + } catch { + // Can't read it here — let the upload run and the server decide. + return true; + } +} + +/** + * Split a selection into what the backend can read and what it will refuse, + * preserving order. Files with a parser pass on their suffix; the rest are + * sniffed, so source, config and log files go through and a video does not. + */ +export async function partitionAttachmentFiles(files: File[]): Promise<{ + supported: File[]; + unsupported: File[]; +}> { + const verdicts = await Promise.all(files.map(isSupportedAttachmentFile)); + const supported: File[] = []; + const unsupported: File[] = []; + files.forEach((file, index) => { + (verdicts[index] ? supported : unsupported).push(file); + }); + return { supported, unsupported }; +} + +/** The ``message`` of a JSON error body, if the server sent one. */ +export function parseUploadErrorMessage(body: string): string | undefined { + if (!body) return undefined; + try { + const parsed = JSON.parse(body) as { message?: unknown }; + return typeof parsed?.message === 'string' && parsed.message + ? parsed.message + : undefined; + } catch { + return undefined; + } +} diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 71a35311..4e8b300d 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -1101,7 +1101,8 @@ "waitingToSend_one": "Will send when {{count}} file finishes processing…", "waitingToSend_other": "Will send when {{count}} files finish processing…", "cancelQueuedSend": "Cancel", - "sendBlockedByFailed": "{{names}} couldn't be processed — remove to send" + "sendBlockedByFailed": "{{names}} couldn't be processed — remove to send", + "unsupportedType": "Not a supported file type ({{extension}})" }, "retry": "Retry", "reasoning": "Reasoning", diff --git a/frontend/src/upload/uploadSlice.ts b/frontend/src/upload/uploadSlice.ts index 468a8353..f512e385 100644 --- a/frontend/src/upload/uploadSlice.ts +++ b/frontend/src/upload/uploadSlice.ts @@ -46,6 +46,8 @@ export interface Attachment { */ attachmentId?: string; token_count?: number; + /** Why a ``failed`` attachment failed, when known (server or client gate). */ + errorMessage?: string; } export type UploadTaskStatus = diff --git a/tests/api/user/attachments/test_routes.py b/tests/api/user/attachments/test_routes.py index d28989af..e320f58d 100644 --- a/tests/api/user/attachments/test_routes.py +++ b/tests/api/user/attachments/test_routes.py @@ -2386,3 +2386,97 @@ class TestRequireLiveSttRedisUnavailable: result = _require_live_stt_redis() assert hasattr(result, "status_code") assert result.status_code == 503 + + +MP4_BYTES = b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomiso2avc1mp41\x00\x00\x00\x08free" + + +class TestStoreAttachmentTypeGate: + """Unparseable uploads are refused before anything is stored or queued.""" + + @patch("application.api.user.tasks.store_attachment.delay") + def test_rejects_unsupported_file_type_before_storage( + self, mock_store_attachment, flask_app + ): + from application.api.user.attachments.routes import StoreAttachment + + app = Flask(__name__) + mock_storage = MagicMock() + with patch("application.api.user.base.storage", mock_storage), app.test_request_context( + "/api/store_attachment", + method="POST", + data={"file": (io.BytesIO(MP4_BYTES), "clip.mp4")}, + content_type="multipart/form-data", + ): + request.decoded_token = {"sub": "test_user"} + response = StoreAttachment().post() + + assert _get_response_status(response) == 400 + body = _get_response_json(response) + assert body["success"] is False + assert body["message"] == "Unsupported file type: .mp4" + assert body["errors"] == [ + {"upload_index": 0, "filename": "clip.mp4", "error": "Unsupported file type: .mp4"} + ] + mock_storage.save_file.assert_not_called() + mock_store_attachment.assert_not_called() + + @patch("application.api.user.tasks.store_attachment.delay") + def test_accepts_a_text_file_with_no_dedicated_parser( + self, mock_store_attachment, flask_app + ): + """A .py or .log is read by the plain-text fallthrough and must stay allowed.""" + from application.api.user.attachments.routes import StoreAttachment + + app = Flask(__name__) + mock_storage = MagicMock() + mock_storage.save_file.return_value = {"storage_type": "local"} + mock_store_attachment.return_value = SimpleNamespace(id="task-py") + + with patch("application.api.user.base.storage", mock_storage), app.test_request_context( + "/api/store_attachment", + method="POST", + data={"file": (io.BytesIO(b"def main():\n return 1\n"), "main.py")}, + content_type="multipart/form-data", + ): + request.decoded_token = {"sub": "test_user"} + response = StoreAttachment().post() + + assert _get_response_status(response) == 200 + assert _get_response_json(response)["task_id"] == "task-py" + assert mock_store_attachment.call_count == 1 + + @patch("application.api.user.tasks.store_attachment.delay") + def test_batch_skips_unsupported_file_and_keeps_the_rest( + self, mock_store_attachment, flask_app + ): + from application.api.user.attachments.routes import StoreAttachment + + app = Flask(__name__) + mock_storage = MagicMock() + mock_storage.save_file.return_value = {"storage_type": "local"} + mock_store_attachment.return_value = SimpleNamespace(id="task-notes") + + with patch("application.api.user.base.storage", mock_storage), app.test_request_context( + "/api/store_attachment", + method="POST", + data={ + "file": [ + (io.BytesIO(MP4_BYTES), "clip.mp4"), + (io.BytesIO(b"hello"), "notes.txt"), + ] + }, + content_type="multipart/form-data", + ): + request.decoded_token = {"sub": "test_user"} + response = StoreAttachment().post() + + assert _get_response_status(response) == 200 + body = _get_response_json(response) + assert [task["upload_index"] for task in body["tasks"]] == [1] + assert body["tasks"][0]["filename"] == "notes.txt" + assert body["errors"] == [ + {"upload_index": 0, "filename": "clip.mp4", "error": "Unsupported file type: .mp4"} + ] + assert mock_storage.save_file.call_count == 1 + assert mock_store_attachment.call_count == 1 diff --git a/tests/parser/file/test_constants.py b/tests/parser/file/test_constants.py new file mode 100644 index 00000000..2bf6f042 --- /dev/null +++ b/tests/parser/file/test_constants.py @@ -0,0 +1,95 @@ +"""Tests for the shared file-extension constants.""" + +import re +from pathlib import Path + +import pytest + +from application.parser.file.bulk import get_default_file_extractor +from application.parser.file.constants import ( + attachment_extension, + has_attachment_parser, + SUPPORTED_ATTACHMENT_EXTENSIONS, + SUPPORTED_SOURCE_EXTENSIONS, +) + +FRONTEND_CONSTANTS = ( + Path(__file__).resolve().parents[3] / "frontend/src/constants/fileUpload.ts" +) + + +@pytest.mark.unit +def test_every_parser_backed_suffix_is_an_allowed_attachment(): + """The allow-list is a hand-kept literal; this is what keeps it honest. + + ``SUPPORTED_ATTACHMENT_EXTENSIONS`` cannot import the extractor (that + would pull docling into the API process), so a parser added to + ``get_default_file_extractor`` would otherwise stay refused as an + attachment — which is how .webp, .tiff, .vtt and .xml were first missed. + """ + parser_suffixes = set(get_default_file_extractor()) + + assert parser_suffixes <= set(SUPPORTED_ATTACHMENT_EXTENSIONS), ( + "parsers exist for suffixes the attachment gate refuses: " + f"{sorted(parser_suffixes - set(SUPPORTED_ATTACHMENT_EXTENSIONS))}" + ) + + +@pytest.mark.unit +def test_frontend_mirrors_the_attachment_allow_list(): + """The composer gates uploads client-side; a divergent list refuses valid files. + + ``ATTACHMENT_PARSER_EXTENSIONS`` in the frontend is a hand-kept copy of + ``SUPPORTED_ATTACHMENT_EXTENSIONS``. Nothing but this test connects them, + and .mdx, .xhtml and .adoc were blocked in the UI while the API accepted + them before it existed. + """ + if not FRONTEND_CONSTANTS.exists(): + pytest.skip("frontend sources not present in this checkout") + + source = FRONTEND_CONSTANTS.read_text(encoding="utf-8") + match = re.search( + r"ATTACHMENT_PARSER_EXTENSIONS:\s*readonly string\[\]\s*=\s*\[(.*?)\]", + source, + re.DOTALL, + ) + assert match, "ATTACHMENT_PARSER_EXTENSIONS not found in fileUpload.ts" + frontend_extensions = set(re.findall(r"'(\.[^']+)'", match.group(1))) + + assert frontend_extensions == set(SUPPORTED_ATTACHMENT_EXTENSIONS), ( + "frontend and backend attachment lists disagree — " + f"backend only: {sorted(set(SUPPORTED_ATTACHMENT_EXTENSIONS) - frontend_extensions)}, " + f"frontend only: {sorted(frontend_extensions - set(SUPPORTED_ATTACHMENT_EXTENSIONS))}" + ) + + +@pytest.mark.unit +def test_attachments_accept_everything_source_ingestion_does(): + assert set(SUPPORTED_SOURCE_EXTENSIONS) <= set(SUPPORTED_ATTACHMENT_EXTENSIONS) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("filename", "expected"), + [ + ("Photo.JPG", ".jpg"), + ("/tmp/dir.d/report.PDF", ".pdf"), + ("archive.tar.gz", ".gz"), + ("Dockerfile", ""), + (".env", ""), + ("trailing.", "."), + (None, ""), + ], +) +def test_attachment_extension(filename, expected): + assert attachment_extension(filename) == expected + + +@pytest.mark.unit +def test_has_attachment_parser_is_case_insensitive_and_excludes_zip(): + assert has_attachment_parser("Report.PDF") + assert has_attachment_parser("scan.WebP") + # No parser: admitted (or not) on content instead, never on the suffix. + assert not has_attachment_parser("archive.zip") + assert not has_attachment_parser("main.py") + assert not has_attachment_parser("clip.mp4") diff --git a/tests/test_upload_limits.py b/tests/test_upload_limits.py index 46818421..acb1fba0 100644 --- a/tests/test_upload_limits.py +++ b/tests/test_upload_limits.py @@ -37,3 +37,126 @@ def test_limited_text_read_decodes_incrementally_and_rejects_overflow(): with pytest.raises(UploadTooLargeError): read_text_upload_limited(_ShortReadStream(b"12345"), max_bytes=4) + + +# --- attachment type gate ------------------------------------------------- +# +# ``SimpleDirectoryReader`` falls through to a plain-text ``open()`` for any +# suffix without a parser. That is what a .py or a .log attachment relies on, +# and it is also how a phone-uploaded video used to be "parsed" into +# megabytes of binary garbage, truncated, and stored with +# ``extraction.status == "ok"``. So a suffix with no parser is admitted on +# content: text in, binary out. + +MP4_HEADER = b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomiso2avc1mp41" + + +@pytest.mark.parametrize( + ("filename", "content"), + [ + ("clip.mp4", MP4_HEADER), + ("movie.MOV", MP4_HEADER), + ("archive.zip", b"PK\x03\x04\x14\x00\x00\x00\x08\x00" + bytes(range(32))), + ("binary", bytes(range(32)) * 8), + ("trailing.", b"\x00\x01\x02\x03"), + ("x.tar.gz", b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03"), + ], +) +def test_enforce_parseable_attachment_rejects_binary_without_a_parser( + filename, content, tmp_path +): + from application.upload_limits import ( + enforce_parseable_attachment, + UnsupportedUploadTypeError, + unsupported_upload_message, + ) + + path = tmp_path / "staged.bin" + path.write_bytes(content) + + with pytest.raises(UnsupportedUploadTypeError) as excinfo: + enforce_parseable_attachment(path, filename) + assert str(excinfo.value) == unsupported_upload_message(filename) + assert str(excinfo.value).startswith("Unsupported file type") + + +@pytest.mark.parametrize( + "filename", + [ + "notes.txt", + "Report.PDF", + "photo.JPG", + "slides.pptx", + "voice.ogg", + "page.xhtml", + "doc.adoc", + "scan.webp", + "fax.tiff", + "subs.vtt", + "feed.xml", + ], +) +def test_enforce_parseable_attachment_accepts_parser_backed_types(filename, tmp_path): + """A parser-backed suffix is admitted on its name — a PDF is binary and parses fine.""" + from application.upload_limits import enforce_parseable_attachment + + path = tmp_path / "staged.bin" + path.write_bytes(MP4_HEADER) + + enforce_parseable_attachment(path, filename) + + +@pytest.mark.parametrize( + "filename", + ["main.py", "server.log", "config.yaml", "query.sql", "Dockerfile", "notes.unknown"], +) +def test_enforce_parseable_attachment_accepts_text_without_a_parser(filename, tmp_path): + """The plain-text fallthrough reads these correctly, so they must stay allowed.""" + from application.upload_limits import enforce_parseable_attachment + + path = tmp_path / "staged.txt" + path.write_text("def main():\n\treturn 'café — ok'\n", encoding="utf-8") + + enforce_parseable_attachment(path, filename) + + +@pytest.mark.parametrize( + ("sample", "expected"), + [ + (b"", True), + (b"plain text\n", True), + ("café — em dash\n".encode(), True), + (b"\x1b[31mred log line\x1b[0m\n", True), + (b"text\x00with nul", False), + (bytes(range(32)) * 4, False), + (b"\x7f\x7f\x7f\x7f" + b"a" * 16, False), + ], +) +def test_looks_like_text(sample, expected): + from application.upload_limits import looks_like_text + + assert looks_like_text(sample) is expected + + +def test_file_looks_like_text_only_samples_the_head(tmp_path): + """Binary past the sampled head is the parser's problem, not the gate's.""" + from application.upload_limits import file_looks_like_text + + path = tmp_path / "staged.log" + path.write_bytes(b"a" * 9000 + b"\x00" * 100) + + assert file_looks_like_text(path) is True + + +def test_file_looks_like_text_allows_an_unreadable_file(tmp_path): + from application.upload_limits import file_looks_like_text + + assert file_looks_like_text(tmp_path / "missing.txt") is True + + +def test_unsupported_upload_message_names_the_extension(): + from application.upload_limits import unsupported_upload_message + + assert unsupported_upload_message("clip.mp4") == "Unsupported file type: .mp4" + assert unsupported_upload_message("Clip.MP4") == "Unsupported file type: .mp4" + assert unsupported_upload_message("binary") == "Unsupported file type: (no extension)" diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index 93e34d97..e5790cae 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -326,3 +326,86 @@ class TestAttachmentZipBombGuard: # BadZipFile → return quietly; the format parser surfaces a clean error. worker._reject_attachment_zip_bomb(str(path)) + + +@pytest.mark.unit +class TestAttachmentTypeGuard: + """A binary with no parser must fail rather than be read as text. + + ``SimpleDirectoryReader`` falls through to a plain-text ``open()`` for a + suffix it has no parser for, which is how a phone-uploaded ``.mp4`` was + once stored as 100k tokens of binary with ``extraction.status == "ok"``. + The same fallthrough is what makes a ``.py`` or ``.log`` attachment work, + so the guard admits those on content. + """ + + def test_binary_without_a_parser_fails_instead_of_being_read_as_text( + self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path + ): + from application import worker + + local_path = tmp_path / "clip.mp4" + local_path.write_bytes(b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomavc1") + + events = [] + fake_storage = MagicMock(name="storage") + fake_storage.process_file.side_effect = lambda path, callback: callback( + str(local_path) + ) + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + monkeypatch.setattr( + worker, + "get_default_file_extractor", + lambda ocr_enabled=False, pdf_text_fast_path=False: {}, + ) + monkeypatch.setattr( + worker, + "publish_user_event", + lambda user, name, payload, **kwargs: events.append((name, payload)), + ) + + file_info = { + "filename": "clip.mp4", + "attachment_id": "507f1f77bcf86cd799439012", + "path": "uploads/user1/attachments/clip.mp4", + "metadata": {"source": "chat"}, + } + + with pytest.raises( + worker.AttachmentRejectedError, match=r"Unsupported file type: \.mp4" + ): + worker.attachment_worker(task_self, file_info, "user1") + + failed = [payload for name, payload in events if name == "attachment.failed"] + assert failed and failed[0]["error"] == "Unsupported file type: .mp4" + + def test_text_without_a_parser_is_parsed( + self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path + ): + from application import worker + + local_path = tmp_path / "server.log" + local_path.write_text("2026-09-02 ERROR boom\n", encoding="utf-8") + + fake_storage = MagicMock(name="storage") + fake_storage.process_file.side_effect = lambda path, callback: callback( + str(local_path) + ) + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + monkeypatch.setattr( + worker, + "get_default_file_extractor", + lambda ocr_enabled=False, pdf_text_fast_path=False: {}, + ) + + file_info = { + "filename": "server.log", + "attachment_id": "507f1f77bcf86cd799439013", + "path": "uploads/user1/attachments/server.log", + "metadata": {"source": "chat"}, + } + + result = worker.attachment_worker(task_self, file_info, "user1") + + assert result["filename"] == "server.log" + assert result["token_count"] > 0 From 6d0319fb63e80070f595061b6f87415e643e99a1 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 3 Sep 2026 09:01:28 +0100 Subject: [PATCH 2/4] feat: better cached token usage logging --- .../versions/0031_token_usage_cache_tokens.py | 35 ++++++ application/llm/base.py | 6 + application/llm/openai.py | 29 ++++- application/storage/db/models.py | 6 + .../storage/db/repositories/token_usage.py | 6 + application/usage.py | 29 ++++- .../handlers/test_round_usage_persistence.py | 8 +- tests/llm/test_base.py | 35 ++++++ tests/llm/test_openai_responses.py | 35 ++++++ .../db/repositories/test_token_usage.py | 29 +++++ tests/storage/db/test_migration_0031.py | 100 ++++++++++++++++ tests/test_usage.py | 112 ++++++++++++++++++ 12 files changed, 423 insertions(+), 7 deletions(-) create mode 100644 application/alembic/versions/0031_token_usage_cache_tokens.py create mode 100644 tests/storage/db/test_migration_0031.py diff --git a/application/alembic/versions/0031_token_usage_cache_tokens.py b/application/alembic/versions/0031_token_usage_cache_tokens.py new file mode 100644 index 00000000..1749b9c2 --- /dev/null +++ b/application/alembic/versions/0031_token_usage_cache_tokens.py @@ -0,0 +1,35 @@ +"""0031 token_usage cache tokens — persist the prompt-cache breakdown. + +Providers report how much of each prompt was served from the prompt cache +(``cached_tokens``) and, on newer OpenAI-family models, how much was written +to it (``cache_write_tokens``). The LLM clients already parsed both and the +usage layer discarded them, so ``token_usage`` could not show a cache hit +rate. These two nullable columns carry the breakdown; NULL means the +provider reported nothing (distinct from 0). ``prompt_tokens`` keeps its +meaning as the provider's total. Idempotent both ways. + +Revision ID: 0031_token_usage_cache_tokens +Revises: 0030_superseded_messages +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0031_token_usage_cache_tokens" +down_revision: Union[str, None] = "0030_superseded_messages" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cached_tokens integer;") + op.execute( + "ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cache_write_tokens integer;" + ) + + +def downgrade() -> None: + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cache_write_tokens;") + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cached_tokens;") diff --git a/application/llm/base.py b/application/llm/base.py index 8e1aabc6..d331ca17 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -710,6 +710,7 @@ class BaseLLM(ABC): completion_tokens, latency_ms, cached_tokens=None, + cache_write_tokens=None, error=None, ): # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired @@ -730,6 +731,8 @@ class BaseLLM(ABC): } if cached_tokens is not None: extra["cached_tokens"] = int(cached_tokens) + if cache_write_tokens is not None: + extra["cache_write_tokens"] = int(cache_write_tokens) if error is not None: extra["error_class"] = type(error).__name__ logging.info("llm_gen_finished", extra=extra) @@ -757,6 +760,7 @@ class BaseLLM(ABC): completion_tokens, latency_ms, cached_tokens=None, + cache_write_tokens=None, error=None, ): # Paired with ``llm_stream_start`` so cost dashboards can sum tokens @@ -774,6 +778,8 @@ class BaseLLM(ABC): } if cached_tokens is not None: extra["cached_tokens"] = int(cached_tokens) + if cache_write_tokens is not None: + extra["cache_write_tokens"] = int(cache_write_tokens) if error is not None: extra["error_class"] = type(error).__name__ logging.info("llm_stream_finished", extra=extra) diff --git a/application/llm/openai.py b/application/llm/openai.py index 1f60e568..65751206 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -1307,17 +1307,36 @@ class OpenAILLM(BaseLLM): output_details = getattr(usage, "completion_tokens_details", None) try: cached = int(getattr(input_details, "cached_tokens", 0) or 0) + written = int(getattr(input_details, "cache_write_tokens", 0) or 0) reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) except (TypeError, ValueError): cached = 0 + written = 0 reasoning = 0 - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + details = self._prompt_cache_details(cached, written) + if details: + result["prompt_tokens_details"] = details if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result self._last_usage_claimed = False + @staticmethod + def _prompt_cache_details(cached: int, written: int) -> dict: + """Build the ``prompt_tokens_details`` breakdown from the two cache bins. + + Only non-zero bins are included so a provider that never reports + caching produces no details (and no misleading zero downstream). + ``cache_write_tokens`` is reported by newer OpenAI-family deployments + that charge for cache writes. + """ + details = {} + if cached: + details["cached_tokens"] = cached + if written: + details["cache_write_tokens"] = written + return details + @staticmethod def _function_call_ids(response): """call_ids of every ``function_call`` item in a Responses output. @@ -1357,9 +1376,11 @@ class OpenAILLM(BaseLLM): "total_tokens": int(getattr(usage, "total_tokens", 0) or prompt + completion), } cached = int(getattr(input_details, "cached_tokens", 0) or 0) + written = int(getattr(input_details, "cache_write_tokens", 0) or 0) reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + details = self._prompt_cache_details(cached, written) + if details: + result["prompt_tokens_details"] = details if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result diff --git a/application/storage/db/models.py b/application/storage/db/models.py index aa80b19b..a65ea3c1 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -235,6 +235,12 @@ token_usage_table = Table( # Added in ``0015_token_usage_model_id``. Canonical model id (catalog # name for built-ins, UUID for BYOM); NULL on un-backfilled rows. Column("model_id", Text), + # Added in ``0031_token_usage_cache_tokens``. Prompt-cache breakdowns of + # ``prompt_tokens`` as reported by the provider (reads; and writes on + # newer OpenAI-family models). NULL = not reported, distinct from 0 = no + # cache activity, so hit-rate queries stay honest across providers. + Column("cached_tokens", Integer), + Column("cache_write_tokens", Integer), ) user_logs_table = Table( diff --git a/application/storage/db/repositories/token_usage.py b/application/storage/db/repositories/token_usage.py index c23ec90b..7292dbd1 100644 --- a/application/storage/db/repositories/token_usage.py +++ b/application/storage/db/repositories/token_usage.py @@ -35,6 +35,8 @@ class TokenUsageRepository: request_id: Optional[str] = None, model_id: Optional[str] = None, timestamp: Optional[datetime] = None, + cached_tokens: Optional[int] = None, + cache_write_tokens: Optional[int] = None, ) -> None: # Attribution guard: the ``token_usage_attribution_chk`` CHECK # constraint requires at least one of ``user_id`` / ``api_key`` @@ -60,12 +62,14 @@ class TokenUsageRepository: INSERT INTO token_usage ( user_id, api_key, agent_id, prompt_tokens, generated_tokens, + cached_tokens, cache_write_tokens, source, request_id, model_id, timestamp ) VALUES ( :user_id, :api_key, CAST(:agent_id AS uuid), :prompt_tokens, :generated_tokens, + :cached_tokens, :cache_write_tokens, :source, :request_id, :model_id, COALESCE(:timestamp, now()) ) """ @@ -76,6 +80,8 @@ class TokenUsageRepository: "agent_id": agent_id_uuid, "prompt_tokens": prompt_tokens, "generated_tokens": generated_tokens, + "cached_tokens": cached_tokens, + "cache_write_tokens": cache_write_tokens, "source": source, "request_id": request_id, "model_id": model_id, diff --git a/application/usage.py b/application/usage.py index a1d31d08..2b10c7aa 100644 --- a/application/usage.py +++ b/application/usage.py @@ -131,6 +131,11 @@ def _persist_call_usage(llm, call_usage): agent_id=str(agent_id) if agent_id else None, prompt_tokens=call_usage["prompt_tokens"], generated_tokens=call_usage["generated_tokens"], + # Present only when the provider reported the breakdown; + # persisted as NULL otherwise so "unknown" never reads as + # "0% cache hits". + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), source=( getattr(llm, "_token_usage_source", None) or "agent_stream" ), @@ -151,6 +156,12 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in ``*_tokens_details`` breakdowns (``cached_tokens``, ``reasoning_tokens``) back out of these bins — that would break parity with what providers bill. + + The prompt-cache sub-bins ARE carried alongside (``cached_tokens``, + ``cache_write_tokens``; Anthropic's ``cache_creation_tokens`` maps to + the latter) so persistence and the finish events can chart them. They + are added only when the provider reported them. The rest of + ``call_usage`` (e.g. ``model``) is preserved rather than replaced. """ reported = getattr(llm, "_last_usage", None) if not isinstance(reported, dict): @@ -174,10 +185,22 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in # Slotted/immutable LLM stand-ins can't record the claim; this # call still gets the provider counts, which is correct for them. pass - return { + merged = { + **call_usage, "prompt_tokens": int(prompt or 0), "generated_tokens": int(completion or 0), } + details = reported.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + written = details.get("cache_write_tokens") + if written is None: + written = details.get("cache_creation_tokens") + if cached is not None: + merged["cached_tokens"] = int(cached or 0) + if written is not None: + merged["cache_write_tokens"] = int(written or 0) + return merged def gen_token_usage(func): @@ -224,6 +247,8 @@ def gen_token_usage(func): prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], latency_ms=int((time.monotonic() - started_at) * 1000), + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, ) except Exception: @@ -272,6 +297,8 @@ def stream_token_usage(func): prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], latency_ms=int((time.monotonic() - started_at) * 1000), + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, ) except Exception: diff --git a/tests/llm/handlers/test_round_usage_persistence.py b/tests/llm/handlers/test_round_usage_persistence.py index 2ada0b9e..f45d418a 100644 --- a/tests/llm/handlers/test_round_usage_persistence.py +++ b/tests/llm/handlers/test_round_usage_persistence.py @@ -179,8 +179,12 @@ class TestPerRoundUsageRows: list(handler.handle_streaming(agent, first, {}, [])) assert len(rows) == 2 - assert rows[0] == {"prompt_tokens": 100, "generated_tokens": 10} - assert rows[1] == {"prompt_tokens": 200, "generated_tokens": 20} + # Provider-reported counts replace the estimates; the rest of the + # call record (e.g. ``model``) rides along, hence a subset check + # rather than dict equality. + counts = lambda row: {k: row[k] for k in ("prompt_tokens", "generated_tokens")} # noqa: E731 + assert counts(rows[0]) == {"prompt_tokens": 100, "generated_tokens": 10} + assert counts(rows[1]) == {"prompt_tokens": 200, "generated_tokens": 20} class TestPreferProviderUsageClaim: diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 518dfa67..67ce9254 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -2347,3 +2347,38 @@ class TestBuildConversationFromMessagesEmpty: result = handler._build_conversation_from_messages(messages) assert result is not None assert len(result.get("queries", [])) >= 1 + + +class TestFinishedLogCacheBins: + """``cached_tokens`` / ``cache_write_tokens`` ride on the finish events + only when a provider reported them, so dashboards can compute a cache + hit rate without a sentinel zero polluting providers that never report.""" + + def test_stream_finished_carries_cache_bins_when_given(self, caplog): + import logging as _logging + + llm = StubLLM() + with caplog.at_level(_logging.INFO, logger="root"): + llm._emit_stream_finished_log( + "m1", + prompt_tokens=100, + completion_tokens=5, + latency_ms=10, + cached_tokens=80, + cache_write_tokens=7, + ) + evt = next(r for r in caplog.records if r.message == "llm_stream_finished") + assert evt.cached_tokens == 80 + assert evt.cache_write_tokens == 7 + + def test_gen_finished_omits_cache_bins_when_none(self, caplog): + import logging as _logging + + llm = StubLLM() + with caplog.at_level(_logging.INFO, logger="root"): + llm._emit_gen_finished_log( + "m1", prompt_tokens=100, completion_tokens=5, latency_ms=10 + ) + evt = next(r for r in caplog.records if r.message == "llm_gen_finished") + assert not hasattr(evt, "cached_tokens") + assert not hasattr(evt, "cache_write_tokens") diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 9ae2e645..ab00a543 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1550,3 +1550,38 @@ def test_responses_gen_stream_retries_without_tools(monkeypatch): assert len(create.calls) == 2 assert "tools" in create.calls[0] assert "tools" not in create.calls[1] + + +@pytest.mark.unit +def test_record_chat_usage_captures_cache_write_bin(monkeypatch): + """Newer OpenAI-family deployments report cache writes next to cache reads.""" + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm._record_chat_usage(_ns( + prompt_tokens=10, + completion_tokens=7, + total_tokens=17, + prompt_tokens_details=_ns(cached_tokens=4, cache_write_tokens=6), + completion_tokens_details=None, + )) + assert llm._last_usage["prompt_tokens_details"] == { + "cached_tokens": 4, + "cache_write_tokens": 6, + } + + +@pytest.mark.unit +def test_record_responses_metadata_captures_cache_write_bin(monkeypatch): + llm = _make_llm(monkeypatch, _responses_caps()) + llm._record_responses_metadata(_ns( + id="resp_1", + usage=_ns( + input_tokens=10, + output_tokens=7, + total_tokens=17, + input_tokens_details=_ns(cached_tokens=0, cache_write_tokens=9), + output_tokens_details=None, + ), + )) + # A pure cache-write turn (first request of a prefix) has reads=0 and + # writes>0; the details must still be recorded. + assert llm._last_usage["prompt_tokens_details"] == {"cache_write_tokens": 9} diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index c4781b4c..681548ef 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -290,3 +290,32 @@ class TestBucketedTotals: ) assert len(rows) == 1 assert rows[0]["prompt_tokens"] == 10 + + +class TestCacheBins: + def _row(self, pg_conn, user_id): + return pg_conn.execute( + text( + "SELECT cached_tokens, cache_write_tokens FROM token_usage " + "WHERE user_id = :u ORDER BY id DESC LIMIT 1" + ), + {"u": user_id}, + ).fetchone() + + def test_persists_cache_bins(self, pg_conn): + repo = TokenUsageRepository(pg_conn) + repo.insert( + user_id="cache-user", + prompt_tokens=1000, + generated_tokens=10, + cached_tokens=800, + cache_write_tokens=100, + ) + assert tuple(self._row(pg_conn, "cache-user")) == (800, 100) + + def test_cache_bins_default_to_null_not_zero(self, pg_conn): + """NULL means "provider did not report"; 0 means "reported no cache + activity". Keeping them distinct is what makes a hit-rate query honest.""" + repo = TokenUsageRepository(pg_conn) + repo.insert(user_id="cache-user-2", prompt_tokens=10, generated_tokens=1) + assert tuple(self._row(pg_conn, "cache-user-2")) == (None, None) diff --git a/tests/storage/db/test_migration_0031.py b/tests/storage/db/test_migration_0031.py new file mode 100644 index 00000000..1836ac3a --- /dev/null +++ b/tests/storage/db/test_migration_0031.py @@ -0,0 +1,100 @@ +"""Migration round-trip test for 0031_token_usage_cache_tokens.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "application" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_heads(url: str) -> list[str]: + out = subprocess.check_output( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + text=True, + ) + return [line for line in out.splitlines() if line.strip()] + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _column_exists(conn, table: str, column: str) -> bool: + row = conn.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'" + ), + {"t": table, "c": column}, + ).fetchone() + return row is not None + + +_0031 = "0031_token_usage_cache_tokens" +_0030 = "0030_superseded_messages" + + +class TestMigration0031RoundTrip: + def test_single_head(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + assert len(_alembic_heads(url)) == 1 + + def test_head_has_cache_columns(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0031 + assert _column_exists(conn, "token_usage", "cached_tokens") + assert _column_exists(conn, "token_usage", "cache_write_tokens") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0030) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0030 + assert not _column_exists(conn, "token_usage", "cached_tokens") + assert not _column_exists(conn, "token_usage", "cache_write_tokens") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0031 + assert _column_exists(conn, "token_usage", "cached_tokens") + + def test_existing_rows_read_null_cache_bins(self, pg_engine): + """Rows written before 0031 must read NULL (unknown), not 0.""" + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0030) + with pg_engine.begin() as conn: + conn.execute( + text( + "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens) " + "VALUES ('u-mig31', 10, 1)" + ) + ) + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + row = conn.execute( + text( + "SELECT cached_tokens, cache_write_tokens FROM token_usage " + "WHERE user_id = 'u-mig31'" + ) + ).fetchone() + assert tuple(row) == (None, None) diff --git a/tests/test_usage.py b/tests/test_usage.py index 0b6053fc..dbc3cdc8 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -600,3 +600,115 @@ class TestCountPromptTokens: ] tokens = _count_prompt_tokens(messages, tools=None) assert tokens > 0 + + +# ── prompt-cache breakdown ──────────────────────────────────────────────────── +# +# Providers report ``cached_tokens`` (and, on newer OpenAI-family models, +# ``cache_write_tokens``) as breakdowns of ``prompt_tokens``. The prompt bin +# stays the provider total (never subtract the details back out); the two +# sub-bins ride alongside so persistence and the finish events can chart a +# cache hit rate. + + +class _ReportingLLM: + decoded_token = {"sub": "user_1"} + user_api_key = None + agent_id = None + token_usage = {"prompt_tokens": 0, "generated_tokens": 0} + + def __init__(self, details): + self._last_usage = { + "prompt_tokens": 1000, + "completion_tokens": 10, + "total_tokens": 1010, + } + if details is not None: + self._last_usage["prompt_tokens_details"] = details + self._last_usage_claimed = False + self.emitted = [] + + def _emit_gen_finished_log(self, model, **kwargs): + self.emitted.append(kwargs) + + +@pytest.mark.unit +def test_prefer_provider_usage_carries_cache_bins(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100}) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert usage["prompt_tokens"] == 1000 + assert usage["generated_tokens"] == 10 + assert usage["cached_tokens"] == 800 + assert usage["cache_write_tokens"] == 100 + + +@pytest.mark.unit +def test_prefer_provider_usage_maps_anthropic_cache_creation_to_writes(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM({"cached_tokens": 3, "cache_creation_tokens": 9}) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert usage["cached_tokens"] == 3 + assert usage["cache_write_tokens"] == 9 + + +@pytest.mark.unit +def test_prefer_provider_usage_without_details_has_no_cache_keys(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM(None) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert "cached_tokens" not in usage + assert "cache_write_tokens" not in usage + + +@pytest.mark.unit +def test_prefer_provider_usage_keeps_the_rest_of_the_call_record(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM(None) + usage = _prefer_provider_usage( + llm, {"prompt_tokens": 1, "generated_tokens": 1, "model": "m"} + ) + assert usage["model"] == "m" + + +@pytest.mark.unit +def test_decorator_persists_and_emits_cache_bins(monkeypatch): + _install_fake_token_repo(monkeypatch) + llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100}) + + @gen_token_usage + def wrapped(self, model, messages, stream, tools, **kwargs): + _ = (model, messages, stream, tools, kwargs) + return "ok" + + wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None) + + row = _FakeTokenUsageRepo.last_instance.inserted[0] + assert row["prompt_tokens"] == 1000 + assert row["cached_tokens"] == 800 + assert row["cache_write_tokens"] == 100 + assert llm.emitted[0]["cached_tokens"] == 800 + assert llm.emitted[0]["cache_write_tokens"] == 100 + + +@pytest.mark.unit +def test_decorator_omits_cache_bins_when_provider_reports_none(monkeypatch): + _install_fake_token_repo(monkeypatch) + llm = _ReportingLLM(None) + + @gen_token_usage + def wrapped(self, model, messages, stream, tools, **kwargs): + _ = (model, messages, stream, tools, kwargs) + return "ok" + + wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None) + + row = _FakeTokenUsageRepo.last_instance.inserted[0] + assert row["cached_tokens"] is None + assert row["cache_write_tokens"] is None + assert llm.emitted[0]["cached_tokens"] is None + assert llm.emitted[0]["cache_write_tokens"] is None From f7cd94668f90f58c2e3ba48913e9e89952f4f4dd Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 3 Sep 2026 09:13:24 +0100 Subject: [PATCH 3/4] fix(attachments): close the .txt gap in the parseability gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the attachment gate. .txt was listed as parser-backed, but it has no parser — it *is* the plain-text fallthrough. That let it skip the content check, so renaming a video to notes.txt walked straight back into the bug the gate exists for (verified: 5132 chars of binary "extracted" and stored). The list is now exactly the file extractor's keys, .txt included in the content check like any other unparsed suffix, and the drift test asserts equality rather than containment. The sniff now recognises a UTF-16/32 BOM as text, so a Notepad "Unicode" .txt is not caught by the NUL-byte rule. The picker's accept filter listed parser-backed suffixes only, hiding .txt, .py and .log — files the gate reads happily — behind "All files". It now carries text/* as well, so it can never be narrower than what the upload accepts. A rejected batch carries one errors entry per file, but the non-200 branch applied the top-level message to every chip, so two files failing for two reasons both reported the first one. Reasons are now matched by upload_index, with the top-level message as fallback. _get_store_attachment_user_error no longer reads str(exc): the unsupported-type message is rebuilt from the filename, so no exception state can reach a response body (CodeQL py/stack-trace-exposure). --- application/api/user/attachments/routes.py | 23 ++++++- application/parser/file/constants.py | 18 +++-- application/upload_limits.py | 21 ++++-- frontend/src/components/MessageInput.tsx | 15 ++++- frontend/src/constants/fileUpload.test.ts | 76 ++++++++++++++++++++-- frontend/src/constants/fileUpload.ts | 61 ++++++++++++++--- tests/api/user/attachments/test_routes.py | 37 +++++++++++ tests/parser/file/test_constants.py | 46 ++++++++----- tests/test_upload_limits.py | 54 ++++++++++++++- 9 files changed, 302 insertions(+), 49 deletions(-) diff --git a/application/api/user/attachments/routes.py b/application/api/user/attachments/routes.py index e8d7491d..3d856cc5 100644 --- a/application/api/user/attachments/routes.py +++ b/application/api/user/attachments/routes.py @@ -42,6 +42,7 @@ from application.upload_limits import ( enforce_parseable_attachment, is_unsupported_upload_message, UnsupportedUploadTypeError, + unsupported_upload_message, upload_limit_message, UploadTooLargeError, ) @@ -128,13 +129,27 @@ def _enforce_uploaded_audio_size_limit(file, filename: str) -> None: enforce_audio_file_size_limit(size_bytes) -def _get_store_attachment_user_error(exc: Exception) -> str: +def _get_store_attachment_user_error(exc: Exception, filename: str | None = None) -> str: + """Map an upload failure to a fixed client-facing message. + + Every branch returns text this module composes itself. Nothing is read off + the exception, so no exception state — message or traceback — can reach a + response body (CodeQL py/stack-trace-exposure). + + Args: + exc: The exception raised while processing one uploaded file. + filename: That file's name, which supplies the extension for the + unsupported-type message. + + Returns: + The message to report for this failure. + """ if isinstance(exc, AudioFileTooLargeError): return build_stt_file_size_limit_message() if isinstance(exc, UploadTooLargeError): return upload_limit_message() if isinstance(exc, UnsupportedUploadTypeError): - return str(exc) + return unsupported_upload_message(filename) return "Failed to process file" @@ -251,7 +266,9 @@ class StoreAttachment(Resource): errors.append({ "upload_index": idx, "filename": file.filename, - "error": _get_store_attachment_user_error(file_err), + "error": _get_store_attachment_user_error( + file_err, file.filename + ), }) if not tasks: diff --git a/application/parser/file/constants.py b/application/parser/file/constants.py index 079a13d2..f8628694 100644 --- a/application/parser/file/constants.py +++ b/application/parser/file/constants.py @@ -28,10 +28,10 @@ SUPPORTED_SOURCE_EXTENSIONS = ( *SUPPORTED_AUDIO_EXTENSIONS, ) -# Suffixes the attachment path has a dedicated parser for — every key of -# ``get_default_file_extractor()`` plus ``.txt``, which needs none. Kept as a -# literal (importing ``bulk`` here would drag docling into the API process), -# with ``tests/parser/file/test_constants.py`` asserting the two agree. +# Suffixes the attachment path has a dedicated parser for — exactly the keys +# of ``get_default_file_extractor()``. Kept as a literal (importing ``bulk`` +# here would drag docling into the API process), with +# ``tests/parser/file/test_constants.py`` asserting the two agree. # # This is *not* the whole attachment allow-list: a suffix with no parser is # read by ``SimpleDirectoryReader``'s plain-text fallthrough, which is right @@ -44,7 +44,7 @@ SUPPORTED_SOURCE_EXTENSIONS = ( # content check like any other binary. # # Mirrored in ``frontend/src/constants/fileUpload.ts``; update both together. -SUPPORTED_ATTACHMENT_EXTENSIONS = frozenset( +ATTACHMENT_PARSER_EXTENSIONS = frozenset( { *SUPPORTED_SOURCE_EXTENSIONS, ".xhtml", @@ -57,6 +57,10 @@ SUPPORTED_ATTACHMENT_EXTENSIONS = frozenset( ".vtt", ".xml", } + # .txt has no parser of its own — it *is* the plain-text fallthrough. It + # must be sniffed like any other unparsed suffix, or renaming a video to + # notes.txt walks straight back into the bug this gate exists for. + - {".txt"} ) @@ -86,6 +90,6 @@ def has_attachment_parser(filename: str | None) -> bool: filename: The upload's filename. Returns: - True when the suffix is in ``SUPPORTED_ATTACHMENT_EXTENSIONS``. + True when the suffix is in ``ATTACHMENT_PARSER_EXTENSIONS``. """ - return attachment_extension(filename) in SUPPORTED_ATTACHMENT_EXTENSIONS + return attachment_extension(filename) in ATTACHMENT_PARSER_EXTENSIONS diff --git a/application/upload_limits.py b/application/upload_limits.py index c5f5c579..3c24c2de 100644 --- a/application/upload_limits.py +++ b/application/upload_limits.py @@ -64,6 +64,14 @@ _TEXT_SNIFF_BYTES = 8192 _MAX_NONTEXT_RATIO = 0.10 # Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. _TEXT_CONTROL_BYTES = frozenset({0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x1B}) +# A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM +# before the NUL test — Notepad's "Unicode" .txt is ordinary text. +_TEXT_BOMS = ( + b"\xef\xbb\xbf", + b"\xff\xfe", + b"\xfe\xff", + b"\x00\x00\xfe\xff", +) def looks_like_text(sample: bytes) -> bool: @@ -74,10 +82,12 @@ def looks_like_text(sample: bytes) -> bool: Returns: False when the sample holds a NUL byte or too many other non-text - control bytes, True otherwise. + control bytes, True otherwise. A Unicode BOM settles it as text. """ if not sample: return True + if sample.startswith(_TEXT_BOMS): + return True if b"\x00" in sample: return False nontext = sum( @@ -110,10 +120,11 @@ def enforce_parseable_attachment( ) -> None: """Reject an attachment that no parser handles and that is not plain text. - Suffixes with a parser (``SUPPORTED_ATTACHMENT_EXTENSIONS``) are admitted - unconditionally — a PDF is binary and parses fine. Everything else has to - read as text, which is what keeps a video or an archive out of the - plain-text fallthrough while leaving source, config and log files in. + Suffixes with a parser (``ATTACHMENT_PARSER_EXTENSIONS``) are admitted + unconditionally — a PDF is binary and parses fine. Everything else, .txt + included, has to read as text: that is what keeps a video or an archive + out of the plain-text fallthrough while leaving source, config and log + files in, whatever the file happens to be named. Args: path: Local path to the staged upload, readable before it is stored. diff --git a/frontend/src/components/MessageInput.tsx b/frontend/src/components/MessageInput.tsx index fac83d66..c472bc58 100644 --- a/frontend/src/components/MessageInput.tsx +++ b/frontend/src/components/MessageInput.tsx @@ -50,6 +50,7 @@ import { AUDIO_FILE_ACCEPT_ATTR, getFileExtension, parseUploadErrorMessage, + parseUploadErrorsByIndex, partitionAttachmentFiles, } from '../constants/fileUpload'; import { UserToolType } from '../settings/types'; @@ -781,12 +782,20 @@ export default function MessageInput({ } } else { console.error('Upload failed', status, xhr.responseText); - const errorMessage = parseUploadErrorMessage(xhr.responseText); - Object.values(indexToUiId).forEach((id) => + // Each file gets its own reason where the server sent one; the + // top-level message is the fallback, not the answer for all of + // them — a batch can fail two files for two different reasons. + const fallbackMessage = parseUploadErrorMessage(xhr.responseText); + const errorsByIndex = parseUploadErrorsByIndex(xhr.responseText); + Object.entries(indexToUiId).forEach(([index, id]) => dispatch( updateAttachment({ id, - updates: { status: 'failed', errorMessage }, + updates: { + status: 'failed', + errorMessage: + errorsByIndex.get(Number(index)) ?? fallbackMessage, + }, }), ), ); diff --git a/frontend/src/constants/fileUpload.test.ts b/frontend/src/constants/fileUpload.test.ts index ff0b9769..02579e76 100644 --- a/frontend/src/constants/fileUpload.test.ts +++ b/frontend/src/constants/fileUpload.test.ts @@ -7,6 +7,7 @@ import { hasAttachmentParser, looksLikeText, parseUploadErrorMessage, + parseUploadErrorsByIndex, partitionAttachmentFiles, } from './fileUpload'; @@ -46,12 +47,12 @@ describe('attachment type gate', () => { expect(hasAttachmentParser(file('main.py'))).toBe(false); expect(hasAttachmentParser(file('archive.zip'))).toBe(false); expect(hasAttachmentParser(file('clip.mp4'))).toBe(false); + // .txt is the plain-text fallthrough itself, not a parser. + expect(hasAttachmentParser(file('notes.txt'))).toBe(false); }); it('mirrors the backend list: no zip, no video, images and markup included', () => { - const listed = ATTACHMENT_FILE_ACCEPT_ATTR.split(','); - expect(listed).toEqual([...ATTACHMENT_PARSER_EXTENSIONS]); - expect(listed).toContain('.pdf'); + expect(ATTACHMENT_PARSER_EXTENSIONS).toContain('.pdf'); // Parser-backed suffixes the first cut of this list missed. for (const ext of [ '.webp', @@ -62,10 +63,22 @@ describe('attachment type gate', () => { '.xml', '.mdx', ]) { + expect(ATTACHMENT_PARSER_EXTENSIONS).toContain(ext); + } + expect(ATTACHMENT_PARSER_EXTENSIONS).not.toContain('.zip'); + expect(ATTACHMENT_PARSER_EXTENSIONS).not.toContain('.mp4'); + expect(ATTACHMENT_PARSER_EXTENSIONS).not.toContain('.txt'); + }); + + it('never hides a file the upload would accept behind the picker filter', () => { + const listed = ATTACHMENT_FILE_ACCEPT_ATTR.split(','); + for (const ext of ATTACHMENT_PARSER_EXTENSIONS) { expect(listed).toContain(ext); } - expect(listed).not.toContain('.zip'); - expect(listed).not.toContain('.mp4'); + // Parserless text (.txt, .py, .log) is accepted on content, so the + // picker must not filter it out. + expect(listed).toContain('.txt'); + expect(listed).toContain('text/*'); }); }); @@ -81,6 +94,14 @@ describe('looksLikeText', () => { ); }); + it('accepts BOM-marked Unicode, which is half NUL bytes but still text', () => { + // UTF-8, UTF-16 LE/BE, UTF-32 BE. + expect(looksLikeText(bytes(0xef, 0xbb, 0xbf, 0x68, 0x69))).toBe(true); + expect(looksLikeText(bytes(0xff, 0xfe, 0x68, 0x00, 0x69, 0x00))).toBe(true); + expect(looksLikeText(bytes(0xfe, 0xff, 0x00, 0x68, 0x00, 0x69))).toBe(true); + expect(looksLikeText(bytes(0x00, 0x00, 0xfe, 0xff, 0x00, 0x68))).toBe(true); + }); + it('rejects a NUL byte and dense control bytes', () => { expect(looksLikeText(MP4_HEADER)).toBe(false); expect(looksLikeText(bytes(0x74, 0x78, 0x00, 0x74))).toBe(false); @@ -104,11 +125,13 @@ describe('partitionAttachmentFiles', () => { it('keeps text files that have no parser — the backend reads them', async () => { const { supported, unsupported } = await partitionAttachmentFiles([ + file('notes.txt', 'plain\n'), file('main.py', 'def main():\n return 1\n'), file('server.log', '2026-09-02 ERROR boom\n'), file('Dockerfile', 'FROM python:3.12\n'), ]); expect(supported.map((f) => f.name)).toEqual([ + 'notes.txt', 'main.py', 'server.log', 'Dockerfile', @@ -116,6 +139,14 @@ describe('partitionAttachmentFiles', () => { expect(unsupported).toEqual([]); }); + it('refuses a binary renamed to .txt — .txt is sniffed like any other suffix', async () => { + const { supported, unsupported } = await partitionAttachmentFiles([ + file('notes.txt', MP4_HEADER), + ]); + expect(supported).toEqual([]); + expect(unsupported.map((f) => f.name)).toEqual(['notes.txt']); + }); + it('admits parser-backed types on their suffix, binary contents and all', async () => { // A PDF is binary; the sniff must never be applied to it. const { supported, unsupported } = await partitionAttachmentFiles([ @@ -155,3 +186,38 @@ describe('parseUploadErrorMessage', () => { expect(parseUploadErrorMessage('')).toBeUndefined(); }); }); + +describe('parseUploadErrorsByIndex', () => { + it('keys each reason by its upload_index', () => { + const byIndex = parseUploadErrorsByIndex( + JSON.stringify({ + success: false, + message: 'Unsupported file type: .mp4', + errors: [ + { + upload_index: 0, + filename: 'clip.mp4', + error: 'Unsupported file type: .mp4', + }, + { + upload_index: 1, + filename: 'a.zip', + error: 'Unsupported file type: .zip', + }, + ], + }), + ); + // Without this, both chips would repeat the first file's reason. + expect(byIndex.get(0)).toBe('Unsupported file type: .mp4'); + expect(byIndex.get(1)).toBe('Unsupported file type: .zip'); + }); + + it('is empty for bodies with no usable errors array', () => { + expect(parseUploadErrorsByIndex('502').size).toBe(0); + expect(parseUploadErrorsByIndex('{"message":"nope"}').size).toBe(0); + expect(parseUploadErrorsByIndex('').size).toBe(0); + expect( + parseUploadErrorsByIndex('{"errors":[{"error":"no index"}]}').size, + ).toBe(0); + }); +}); diff --git a/frontend/src/constants/fileUpload.ts b/frontend/src/constants/fileUpload.ts index b8fd7c0d..d68b2aa3 100644 --- a/frontend/src/constants/fileUpload.ts +++ b/frontend/src/constants/fileUpload.ts @@ -83,20 +83,20 @@ export const SOURCE_FILE_TREE_ACCEPT_ATTR = [ /** * Chat-attachment suffixes with a dedicated parser. Mirrors the backend's - * `SUPPORTED_ATTACHMENT_EXTENSIONS` (application/parser/file/constants.py) — + * `ATTACHMENT_PARSER_EXTENSIONS` (application/parser/file/constants.py) — * update both together. Zip is absent: source ingestion extracts archives, * the attachment path does not. * - * Not the whole allow-list. A suffix that isn't here (.py, .log, .yaml) is - * read by the backend's plain-text fallthrough, so it is judged on content - * by `partitionAttachmentFiles` instead, exactly as the server does. + * Not the whole allow-list, and `.txt` is deliberately not here: a suffix + * that isn't listed (.txt, .py, .log, .yaml) is read by the backend's + * plain-text fallthrough, so it is judged on content by + * `partitionAttachmentFiles` instead, exactly as the server does. */ export const ATTACHMENT_PARSER_EXTENSIONS: readonly string[] = [ '.rst', '.md', '.mdx', '.pdf', - '.txt', '.docx', '.csv', '.epub', @@ -123,9 +123,17 @@ export const ATTACHMENT_PARSER_EXTENSIONS: readonly string[] = [ '.webm', ]; -/** Picker filter for the Attach button. A hint only — pickers may ignore it. */ -export const ATTACHMENT_FILE_ACCEPT_ATTR = - ATTACHMENT_PARSER_EXTENSIONS.join(','); +/** + * Picker filter for the Attach button. A hint only — pickers may ignore it, + * and it must never be narrower than what the upload accepts: `text/*` (plus + * `.txt` explicitly) keeps parserless text files such as .py and .log + * selectable, since the gate reads those happily. + */ +export const ATTACHMENT_FILE_ACCEPT_ATTR = [ + ...ATTACHMENT_PARSER_EXTENSIONS, + '.txt', + 'text/*', +].join(','); /** Lower-cased last extension including the dot, or '' (dotfiles have none). */ export function getFileExtension(name: string): string { @@ -152,10 +160,20 @@ const TEXT_SNIFF_BYTES = 8192; const MAX_NONTEXT_RATIO = 0.1; // Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. const TEXT_CONTROL_BYTES = new Set([0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x1b]); +// A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM +// before the NUL test — Notepad's "Unicode" .txt is ordinary text. +const TEXT_BOMS = [ + [0xef, 0xbb, 0xbf], + [0xff, 0xfe], + [0xfe, 0xff], + [0x00, 0x00, 0xfe, 0xff], +]; /** Whether a leading byte sample reads as text rather than binary. */ export function looksLikeText(sample: Uint8Array): boolean { if (sample.length === 0) return true; + if (TEXT_BOMS.some((bom) => bom.every((byte, i) => sample[i] === byte))) + return true; let nontext = 0; for (const byte of sample) { if (byte === 0x00) return false; @@ -206,3 +224,30 @@ export function parseUploadErrorMessage(body: string): string | undefined { return undefined; } } + +/** + * Per-file reasons from an error body, keyed by `upload_index`. A rejected + * batch carries one `errors` entry per file, so each chip can say why it + * failed instead of every chip repeating the first file's reason. + */ +export function parseUploadErrorsByIndex(body: string): Map { + const byIndex = new Map(); + if (!body) return byIndex; + try { + const parsed = JSON.parse(body) as { errors?: unknown }; + if (!Array.isArray(parsed?.errors)) return byIndex; + for (const entry of parsed.errors as { + upload_index?: unknown; + error?: unknown; + }[]) { + if ( + typeof entry?.upload_index === 'number' && + typeof entry.error === 'string' + ) + byIndex.set(entry.upload_index, entry.error); + } + } catch { + // Not JSON (a proxy's HTML 502, say) — the caller falls back. + } + return byIndex; +} diff --git a/tests/api/user/attachments/test_routes.py b/tests/api/user/attachments/test_routes.py index e320f58d..a32fa36d 100644 --- a/tests/api/user/attachments/test_routes.py +++ b/tests/api/user/attachments/test_routes.py @@ -2369,6 +2369,20 @@ class TestGetStoreAttachmentUserError: msg = _get_store_attachment_user_error(RuntimeError("oops")) assert msg == "Failed to process file" + @pytest.mark.unit + def test_unsupported_type_message_is_rebuilt_from_the_filename(self): + """Nothing is read off the exception, so its text cannot reach a response.""" + from application.api.user.attachments.routes import ( + _get_store_attachment_user_error, + ) + from application.upload_limits import UnsupportedUploadTypeError + + err = UnsupportedUploadTypeError("internals: /srv/app/tmp/staged-42") + msg = _get_store_attachment_user_error(err, "clip.mp4") + + assert msg == "Unsupported file type: .mp4" + assert "/srv/app" not in msg + class TestRequireLiveSttRedisUnavailable: """Cover lines 99-102: Redis unavailable returns 503.""" @@ -2421,6 +2435,29 @@ class TestStoreAttachmentTypeGate: mock_storage.save_file.assert_not_called() mock_store_attachment.assert_not_called() + @patch("application.api.user.tasks.store_attachment.delay") + def test_rejects_a_binary_renamed_to_a_text_suffix( + self, mock_store_attachment, flask_app + ): + """.txt has no parser, so it is judged on content like any other suffix.""" + from application.api.user.attachments.routes import StoreAttachment + + app = Flask(__name__) + mock_storage = MagicMock() + with patch("application.api.user.base.storage", mock_storage), app.test_request_context( + "/api/store_attachment", + method="POST", + data={"file": (io.BytesIO(MP4_BYTES), "notes.txt")}, + content_type="multipart/form-data", + ): + request.decoded_token = {"sub": "test_user"} + response = StoreAttachment().post() + + assert _get_response_status(response) == 400 + assert _get_response_json(response)["message"] == "Unsupported file type: .txt" + mock_storage.save_file.assert_not_called() + mock_store_attachment.assert_not_called() + @patch("application.api.user.tasks.store_attachment.delay") def test_accepts_a_text_file_with_no_dedicated_parser( self, mock_store_attachment, flask_app diff --git a/tests/parser/file/test_constants.py b/tests/parser/file/test_constants.py index 2bf6f042..f0cf3919 100644 --- a/tests/parser/file/test_constants.py +++ b/tests/parser/file/test_constants.py @@ -1,15 +1,16 @@ """Tests for the shared file-extension constants.""" import re +from importlib.util import find_spec from pathlib import Path import pytest from application.parser.file.bulk import get_default_file_extractor from application.parser.file.constants import ( + ATTACHMENT_PARSER_EXTENSIONS, attachment_extension, has_attachment_parser, - SUPPORTED_ATTACHMENT_EXTENSIONS, SUPPORTED_SOURCE_EXTENSIONS, ) @@ -19,30 +20,38 @@ FRONTEND_CONSTANTS = ( @pytest.mark.unit -def test_every_parser_backed_suffix_is_an_allowed_attachment(): - """The allow-list is a hand-kept literal; this is what keeps it honest. +def test_parser_extensions_match_the_extractor(): + """The list is a hand-kept literal; this is what keeps it honest. - ``SUPPORTED_ATTACHMENT_EXTENSIONS`` cannot import the extractor (that - would pull docling into the API process), so a parser added to + ``ATTACHMENT_PARSER_EXTENSIONS`` cannot import the extractor (that would + pull docling into the API process), so a parser added to ``get_default_file_extractor`` would otherwise stay refused as an attachment — which is how .webp, .tiff, .vtt and .xml were first missed. + A suffix listed here but *without* a parser is the opposite failure: it + skips the content check and reaches the plain-text fallthrough, which is + how a binary named notes.txt got through. """ parser_suffixes = set(get_default_file_extractor()) - assert parser_suffixes <= set(SUPPORTED_ATTACHMENT_EXTENSIONS), ( + assert parser_suffixes <= set(ATTACHMENT_PARSER_EXTENSIONS), ( "parsers exist for suffixes the attachment gate refuses: " - f"{sorted(parser_suffixes - set(SUPPORTED_ATTACHMENT_EXTENSIONS))}" + f"{sorted(parser_suffixes - set(ATTACHMENT_PARSER_EXTENSIONS))}" ) + if find_spec("docling") is not None: + assert parser_suffixes == set(ATTACHMENT_PARSER_EXTENSIONS), ( + "listed as parser-backed but nothing parses them: " + f"{sorted(set(ATTACHMENT_PARSER_EXTENSIONS) - parser_suffixes)}" + ) @pytest.mark.unit -def test_frontend_mirrors_the_attachment_allow_list(): +def test_frontend_mirrors_the_parser_extensions(): """The composer gates uploads client-side; a divergent list refuses valid files. ``ATTACHMENT_PARSER_EXTENSIONS`` in the frontend is a hand-kept copy of - ``SUPPORTED_ATTACHMENT_EXTENSIONS``. Nothing but this test connects them, - and .mdx, .xhtml and .adoc were blocked in the UI while the API accepted - them before it existed. + the backend's. Nothing but this test connects them, and .mdx, .xhtml and + .adoc were blocked in the UI while the API accepted them before it + existed. """ if not FRONTEND_CONSTANTS.exists(): pytest.skip("frontend sources not present in this checkout") @@ -56,16 +65,19 @@ def test_frontend_mirrors_the_attachment_allow_list(): assert match, "ATTACHMENT_PARSER_EXTENSIONS not found in fileUpload.ts" frontend_extensions = set(re.findall(r"'(\.[^']+)'", match.group(1))) - assert frontend_extensions == set(SUPPORTED_ATTACHMENT_EXTENSIONS), ( + assert frontend_extensions == set(ATTACHMENT_PARSER_EXTENSIONS), ( "frontend and backend attachment lists disagree — " - f"backend only: {sorted(set(SUPPORTED_ATTACHMENT_EXTENSIONS) - frontend_extensions)}, " - f"frontend only: {sorted(frontend_extensions - set(SUPPORTED_ATTACHMENT_EXTENSIONS))}" + f"backend only: {sorted(set(ATTACHMENT_PARSER_EXTENSIONS) - frontend_extensions)}, " + f"frontend only: {sorted(frontend_extensions - set(ATTACHMENT_PARSER_EXTENSIONS))}" ) @pytest.mark.unit -def test_attachments_accept_everything_source_ingestion_does(): - assert set(SUPPORTED_SOURCE_EXTENSIONS) <= set(SUPPORTED_ATTACHMENT_EXTENSIONS) +def test_source_ingestion_types_are_parser_backed_apart_from_text(): + """Source ingestion's own list must not smuggle a parserless suffix past the sniff.""" + assert set(SUPPORTED_SOURCE_EXTENSIONS) - set(ATTACHMENT_PARSER_EXTENSIONS) == { + ".txt" + } @pytest.mark.unit @@ -93,3 +105,5 @@ def test_has_attachment_parser_is_case_insensitive_and_excludes_zip(): assert not has_attachment_parser("archive.zip") assert not has_attachment_parser("main.py") assert not has_attachment_parser("clip.mp4") + # .txt is the plain-text fallthrough itself, not a parser. + assert not has_attachment_parser("notes.txt") diff --git a/tests/test_upload_limits.py b/tests/test_upload_limits.py index acb1fba0..87c868f3 100644 --- a/tests/test_upload_limits.py +++ b/tests/test_upload_limits.py @@ -1,5 +1,6 @@ """Tests for bounded user-upload stream helpers.""" +import codecs import io import pytest @@ -80,10 +81,49 @@ def test_enforce_parseable_attachment_rejects_binary_without_a_parser( assert str(excinfo.value).startswith("Unsupported file type") +def test_enforce_parseable_attachment_rejects_binary_named_as_text(tmp_path): + """.txt has no parser — it *is* the fallthrough — so it is sniffed like any suffix. + + Renaming a video to notes.txt would otherwise walk straight back into the + bug this gate exists for. + """ + from application.upload_limits import ( + enforce_parseable_attachment, + UnsupportedUploadTypeError, + ) + + path = tmp_path / "notes.txt" + path.write_bytes(MP4_HEADER + bytes(range(256)) * 4) + + with pytest.raises(UnsupportedUploadTypeError) as excinfo: + enforce_parseable_attachment(path, "notes.txt") + assert str(excinfo.value) == "Unsupported file type: .txt" + + +@pytest.mark.parametrize( + "content", + [ + codecs.BOM_UTF8 + "hello — Unicode\n".encode(), + codecs.BOM_UTF16_LE + "hello\n".encode("utf-16-le"), + codecs.BOM_UTF16_BE + "hello\n".encode("utf-16-be"), + codecs.BOM_UTF32_BE + "hello\n".encode("utf-32-be"), + ], +) +def test_enforce_parseable_attachment_accepts_bom_marked_unicode_text( + content, tmp_path +): + """A UTF-16 .txt is half NUL bytes and still ordinary text — the BOM says so.""" + from application.upload_limits import enforce_parseable_attachment + + path = tmp_path / "notes.txt" + path.write_bytes(content) + + enforce_parseable_attachment(path, "notes.txt") + + @pytest.mark.parametrize( "filename", [ - "notes.txt", "Report.PDF", "photo.JPG", "slides.pptx", @@ -108,7 +148,15 @@ def test_enforce_parseable_attachment_accepts_parser_backed_types(filename, tmp_ @pytest.mark.parametrize( "filename", - ["main.py", "server.log", "config.yaml", "query.sql", "Dockerfile", "notes.unknown"], + [ + "notes.txt", + "main.py", + "server.log", + "config.yaml", + "query.sql", + "Dockerfile", + "notes.unknown", + ], ) def test_enforce_parseable_attachment_accepts_text_without_a_parser(filename, tmp_path): """The plain-text fallthrough reads these correctly, so they must stay allowed.""" @@ -127,6 +175,8 @@ def test_enforce_parseable_attachment_accepts_text_without_a_parser(filename, tm (b"plain text\n", True), ("café — em dash\n".encode(), True), (b"\x1b[31mred log line\x1b[0m\n", True), + (codecs.BOM_UTF16_LE + "hi\n".encode("utf-16-le"), True), + (codecs.BOM_UTF8 + b"hi\n", True), (b"text\x00with nul", False), (bytes(range(32)) * 4, False), (b"\x7f\x7f\x7f\x7f" + b"a" * 16, False), From 9db11f18f5fddad36f49a073d25f728cd0ee3e7e Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 3 Sep 2026 09:46:55 +0100 Subject: [PATCH 4/4] fix(attachments,usage): validate content behind a BOM, judge by the live parser table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups. A BOM told the sniff which encoding to read, but was also taken as the verdict: three prepended bytes let any binary through, including as notes.txt. A BOM now only selects the test — UTF-8 falls through to the byte rules on the remainder, UTF-16/32 decode and judge the characters (NUL, unprintable, or replacement chars from bytes the decoder could not read). Real Notepad-Unicode text still passes, mp4-behind-a-BOM does not, in either language. The gate treated the full parser table as a given, but without docling the fallback extractor has no .tif/.tiff/.bmp/.webp/.vtt/.xml handler, so those suffixes skipped the content check and reached the plain-text fallthrough — the original bug, one install away. The worker now passes the keys of the extractor it actually built, making the second gate stricter than the route's static one rather than a copy of it. Cache bins: extraction coerced a missing value to 0 and only non-zero bins were recorded, so a provider reporting cached_tokens=0 persisted as NULL — indistinguishable from "not reported", and OpenAI reports exactly that on every uncached request. Bins are now carried as Optional and recorded when not None, which is what the nullable columns and the NULL-means-unknown comment already assumed. Anthropic's cache_read/cache_creation bins had the same shape and are fixed alongside; the int-or-None coercion is shared in llm/base.py. --- application/llm/anthropic.py | 16 ++-- application/llm/base.py | 23 ++++++ application/llm/openai.py | 42 ++++++----- application/upload_limits.py | 91 ++++++++++++++++++----- application/worker.py | 16 ++-- frontend/src/constants/fileUpload.test.ts | 21 +++++- frontend/src/constants/fileUpload.ts | 79 +++++++++++++++++--- tests/llm/test_openai_responses.py | 25 ++++++- tests/test_upload_limits.py | 47 ++++++++++++ tests/worker/test_attachment_worker.py | 47 ++++++++++++ 10 files changed, 341 insertions(+), 66 deletions(-) diff --git a/application/llm/anthropic.py b/application/llm/anthropic.py index 7944dff6..f26a00ac 100644 --- a/application/llm/anthropic.py +++ b/application/llm/anthropic.py @@ -6,7 +6,7 @@ from typing import Any, Dict, Generator, List, Optional, Tuple from anthropic import Anthropic from application.core.settings import settings -from application.llm.base import BaseLLM +from application.llm.base import BaseLLM, optional_int from application.storage.storage_creator import StorageCreator logger = logging.getLogger(__name__) @@ -387,8 +387,12 @@ class AnthropicLLM(BaseLLM): # message_start are reused below instead. if not output_only: base_input = int(getattr(usage, "input_tokens", 0) or 0) - created = int(getattr(usage, "cache_creation_input_tokens", 0) or 0) - cached = int(getattr(usage, "cache_read_input_tokens", 0) or 0) + # None (unreported) and 0 (reported, no caching) must stay + # distinguishable — see ``optional_int``. + created = optional_int( + getattr(usage, "cache_creation_input_tokens", None) + ) + cached = optional_int(getattr(usage, "cache_read_input_tokens", None)) except (TypeError, ValueError): return @@ -398,11 +402,11 @@ class AnthropicLLM(BaseLLM): prompt = int(previous.get("prompt_tokens", 0) or 0) details = previous.get("prompt_tokens_details") else: - prompt = base_input + created + cached + prompt = base_input + (created or 0) + (cached or 0) details = {} - if cached: + if cached is not None: details["cached_tokens"] = cached - if created: + if created is not None: details["cache_creation_tokens"] = created details = details or None diff --git a/application/llm/base.py b/application/llm/base.py index d331ca17..2afa2f2a 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -33,6 +33,29 @@ _STREAM_RETRYABLE_TRANSPORT_ERRORS = ( ) +def optional_int(value) -> Optional[int]: + """Coerce a provider-reported count to int, keeping "unreported" as None. + + Cache bins persist as NULL when a provider says nothing and as 0 when it + reports zero, so the two must stay distinguishable all the way from the + usage object: providers report ``cached_tokens: 0`` on every uncached + request, and folding that into "unknown" would file the ordinary case as + no-data. + + Args: + value: The raw attribute off a provider usage object. + + Returns: + The integer value, or None when absent or not a number. + """ + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + class BaseLLM(ABC): # Stamped onto the ``llm_stream_start`` event so dashboards can group # calls by vendor. Subclasses override. diff --git a/application/llm/openai.py b/application/llm/openai.py index 65751206..6fcc8bc3 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -9,7 +9,7 @@ from typing import Any, Callable from openai import BadRequestError, OpenAI from application.core.settings import settings -from application.llm.base import BaseLLM +from application.llm.base import BaseLLM, optional_int from application.storage.storage_creator import StorageCreator # Placeholder sent to OpenAI-compatible backends that require no credentials. @@ -1305,14 +1305,11 @@ class OpenAILLM(BaseLLM): } input_details = getattr(usage, "prompt_tokens_details", None) output_details = getattr(usage, "completion_tokens_details", None) - try: - cached = int(getattr(input_details, "cached_tokens", 0) or 0) - written = int(getattr(input_details, "cache_write_tokens", 0) or 0) - reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) - except (TypeError, ValueError): - cached = 0 - written = 0 - reasoning = 0 + cached = optional_int(getattr(input_details, "cached_tokens", None)) + written = optional_int(getattr(input_details, "cache_write_tokens", None)) + reasoning = optional_int( + getattr(output_details, "reasoning_tokens", None) + ) details = self._prompt_cache_details(cached, written) if details: result["prompt_tokens_details"] = details @@ -1322,18 +1319,21 @@ class OpenAILLM(BaseLLM): self._last_usage_claimed = False @staticmethod - def _prompt_cache_details(cached: int, written: int) -> dict: + def _prompt_cache_details(cached: int | None, written: int | None) -> dict: """Build the ``prompt_tokens_details`` breakdown from the two cache bins. - Only non-zero bins are included so a provider that never reports - caching produces no details (and no misleading zero downstream). - ``cache_write_tokens`` is reported by newer OpenAI-family deployments - that charge for cache writes. + A bin is included whenever the provider reported it, zero included: + downstream, an absent bin persists as NULL ("we don't know") and a + reported zero as 0 ("no cache hits"), and OpenAI reports + ``cached_tokens: 0`` on every uncached request. Collapsing the two + would file every ordinary request as unknown. ``cache_write_tokens`` + is reported by newer OpenAI-family deployments that charge for cache + writes. """ details = {} - if cached: + if cached is not None: details["cached_tokens"] = cached - if written: + if written is not None: details["cache_write_tokens"] = written return details @@ -1375,9 +1375,13 @@ class OpenAILLM(BaseLLM): "completion_tokens": completion, "total_tokens": int(getattr(usage, "total_tokens", 0) or prompt + completion), } - cached = int(getattr(input_details, "cached_tokens", 0) or 0) - written = int(getattr(input_details, "cache_write_tokens", 0) or 0) - reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) + cached = optional_int(getattr(input_details, "cached_tokens", None)) + written = optional_int( + getattr(input_details, "cache_write_tokens", None) + ) + reasoning = optional_int( + getattr(output_details, "reasoning_tokens", None) + ) details = self._prompt_cache_details(cached, written) if details: result["prompt_tokens_details"] = details diff --git a/application/upload_limits.py b/application/upload_limits.py index 3c24c2de..7ea3b0a8 100644 --- a/application/upload_limits.py +++ b/application/upload_limits.py @@ -2,10 +2,11 @@ from __future__ import annotations +import codecs import io import os from contextlib import suppress -from typing import BinaryIO +from typing import BinaryIO, Container from application.core.settings import settings from application.parser.file.constants import ( @@ -64,16 +65,48 @@ _TEXT_SNIFF_BYTES = 8192 _MAX_NONTEXT_RATIO = 0.10 # Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. _TEXT_CONTROL_BYTES = frozenset({0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x1B}) -# A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM -# before the NUL test — Notepad's "Unicode" .txt is ordinary text. +_TEXT_CONTROL_CHARS = frozenset("\t\n\v\f\r\x1b") +# A UTF-16/32 file is half NUL bytes, so the byte rules below would reject it. +# Its BOM says which encoding to read it as; the decoded characters are then +# judged instead. Longest BOM first — UTF-32-LE starts with the UTF-16-LE one. +# A BOM is never a verdict on its own: three bytes must not buy a video a pass. _TEXT_BOMS = ( - b"\xef\xbb\xbf", - b"\xff\xfe", - b"\xfe\xff", - b"\x00\x00\xfe\xff", + (codecs.BOM_UTF32_LE, "utf-32-le"), + (codecs.BOM_UTF32_BE, "utf-32-be"), + (codecs.BOM_UTF8, None), + (codecs.BOM_UTF16_LE, "utf-16-le"), + (codecs.BOM_UTF16_BE, "utf-16-be"), ) +def _decoded_looks_like_text(text: str) -> bool: + """Return whether decoded characters read as text. + + Args: + text: Characters decoded from a BOM-marked sample. + + Returns: + False when a NUL character appears — the byte-level rule, one level + up, since no text holds one — or when too many characters are + unprintable: control, unassigned, private-use or surrogate, which is + what binary decodes into. Replacement characters count as unprintable, + so bytes the decoder could not read are evidence rather than silently + dropped (a truncated character at the sample boundary is one of + thousands and cannot swing the ratio). + """ + if not text: + return True + if "\x00" in text: + return False + nontext = sum( + 1 + for char in text + if char == "�" + or (not char.isprintable() and char not in _TEXT_CONTROL_CHARS) + ) + return nontext / len(text) <= _MAX_NONTEXT_RATIO + + def looks_like_text(sample: bytes) -> bool: """Return whether a leading byte sample reads as text rather than binary. @@ -82,20 +115,30 @@ def looks_like_text(sample: bytes) -> bool: Returns: False when the sample holds a NUL byte or too many other non-text - control bytes, True otherwise. A Unicode BOM settles it as text. + control bytes, True otherwise. A UTF-16/32 BOM switches the test to + the decoded characters; the content is still what decides. """ if not sample: return True - if sample.startswith(_TEXT_BOMS): + body = sample + for bom, encoding in _TEXT_BOMS: + if not sample.startswith(bom): + continue + body = sample[len(bom) :] + if encoding is not None: + return _decoded_looks_like_text(body.decode(encoding, errors="replace")) + # UTF-8 BOM: the byte rules still apply to everything after it. + break + if not body: return True - if b"\x00" in sample: + if b"\x00" in body: return False nontext = sum( 1 - for byte in sample + for byte in body if (byte < 0x20 and byte not in _TEXT_CONTROL_BYTES) or byte == 0x7F ) - return nontext / len(sample) <= _MAX_NONTEXT_RATIO + return nontext / len(body) <= _MAX_NONTEXT_RATIO def file_looks_like_text(path: str | os.PathLike[str]) -> bool: @@ -116,25 +159,35 @@ def file_looks_like_text(path: str | os.PathLike[str]) -> bool: def enforce_parseable_attachment( - path: str | os.PathLike[str], filename: str | None + path: str | os.PathLike[str], + filename: str | None, + parser_extensions: Container[str] | None = None, ) -> None: """Reject an attachment that no parser handles and that is not plain text. - Suffixes with a parser (``ATTACHMENT_PARSER_EXTENSIONS``) are admitted - unconditionally — a PDF is binary and parses fine. Everything else, .txt - included, has to read as text: that is what keeps a video or an archive - out of the plain-text fallthrough while leaving source, config and log - files in, whatever the file happens to be named. + Suffixes with a parser are admitted unconditionally — a PDF is binary and + parses fine. Everything else, .txt included, has to read as text: that is + what keeps a video or an archive out of the plain-text fallthrough while + leaving source, config and log files in, whatever the file is named. Args: path: Local path to the staged upload, readable before it is stored. filename: The upload's filename, used for the suffix and the message. + parser_extensions: Suffixes to treat as parser-backed. Defaults to + ``ATTACHMENT_PARSER_EXTENSIONS``, which assumes the full parser + table; callers holding the live extractor (the worker) pass its + keys, so a docling-less install does not admit a .webp on trust + and then read it as text. Raises: UnsupportedUploadTypeError: When the file has no parser and its contents are binary. """ - if has_attachment_parser(filename): + if parser_extensions is None: + has_parser = has_attachment_parser(filename) + else: + has_parser = attachment_extension(filename) in parser_extensions + if has_parser: return if file_looks_like_text(path): return diff --git a/application/worker.py b/application/worker.py index a1507239..168985c6 100755 --- a/application/worker.py +++ b/application/worker.py @@ -1558,22 +1558,26 @@ class AttachmentRejectedError(Exception): """ -def _reject_unparseable_attachment(local_path: str, filename: str) -> None: +def _reject_unparseable_attachment( + local_path: str, filename: str, parser_extensions +) -> None: """Reject an attachment with no parser whose contents are binary. - Defence in depth behind the route's gate: ``SimpleDirectoryReader`` opens - any suffix it has no parser for as plain text, so a binary must fail here, - visibly, rather than "succeed" as garbage. + Defence in depth behind the route's gate, and stricter than it: this runs + against the parser table actually loaded, so a suffix the route trusted + (.webp with docling absent, say) is still refused rather than opened as + plain text by ``SimpleDirectoryReader`` and stored as garbage. Args: local_path: Filesystem path of the attachment about to be parsed. filename: The upload's original filename, which carries the suffix. + parser_extensions: Suffixes the live ``file_extractor`` can parse. Raises: AttachmentRejectedError: If the file has no parser and is not text. """ try: - enforce_parseable_attachment(local_path, filename) + enforce_parseable_attachment(local_path, filename, parser_extensions) except UnsupportedUploadTypeError as exc: raise AttachmentRejectedError(str(exc)) from exc @@ -1775,7 +1779,7 @@ def attachment_worker(self, file_info, user): parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader" def _parse_local_file(local_path: str, **kwargs) -> Document: - _reject_unparseable_attachment(local_path, filename) + _reject_unparseable_attachment(local_path, filename, set(file_extractor)) _reject_attachment_zip_bomb(local_path) parse_path, is_temp_copy = _bounded_attachment_copy(local_path) try: diff --git a/frontend/src/constants/fileUpload.test.ts b/frontend/src/constants/fileUpload.test.ts index 02579e76..c3fa4839 100644 --- a/frontend/src/constants/fileUpload.test.ts +++ b/frontend/src/constants/fileUpload.test.ts @@ -95,13 +95,28 @@ describe('looksLikeText', () => { }); it('accepts BOM-marked Unicode, which is half NUL bytes but still text', () => { - // UTF-8, UTF-16 LE/BE, UTF-32 BE. + // UTF-8, UTF-16 LE/BE, and UTF-32 BE (no TextDecoder — left to the server). expect(looksLikeText(bytes(0xef, 0xbb, 0xbf, 0x68, 0x69))).toBe(true); expect(looksLikeText(bytes(0xff, 0xfe, 0x68, 0x00, 0x69, 0x00))).toBe(true); expect(looksLikeText(bytes(0xfe, 0xff, 0x00, 0x68, 0x00, 0x69))).toBe(true); expect(looksLikeText(bytes(0x00, 0x00, 0xfe, 0xff, 0x00, 0x68))).toBe(true); }); + it('still rejects binary behind a BOM — three bytes buy nothing', () => { + // UTF-8 BOM + MP4 header: the byte rules apply to what follows it. + expect(looksLikeText(bytes(0xef, 0xbb, 0xbf, ...MP4_HEADER))).toBe(false); + // UTF-16 LE BOM + control-only code points. + expect( + looksLikeText( + bytes(0xff, 0xfe, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04, 0x00), + ), + ).toBe(false); + // UTF-16 BE BOM + bytes that decode to NUL characters. + expect(looksLikeText(bytes(0xfe, 0xff, 0x00, 0x00, 0x00, 0x41))).toBe( + false, + ); + }); + it('rejects a NUL byte and dense control bytes', () => { expect(looksLikeText(MP4_HEADER)).toBe(false); expect(looksLikeText(bytes(0x74, 0x78, 0x00, 0x74))).toBe(false); @@ -142,9 +157,11 @@ describe('partitionAttachmentFiles', () => { it('refuses a binary renamed to .txt — .txt is sniffed like any other suffix', async () => { const { supported, unsupported } = await partitionAttachmentFiles([ file('notes.txt', MP4_HEADER), + // A BOM in front of it changes nothing. + file('bom.txt', bytes(0xef, 0xbb, 0xbf, ...MP4_HEADER)), ]); expect(supported).toEqual([]); - expect(unsupported.map((f) => f.name)).toEqual(['notes.txt']); + expect(unsupported.map((f) => f.name)).toEqual(['notes.txt', 'bom.txt']); }); it('admits parser-backed types on their suffix, binary contents and all', async () => { diff --git a/frontend/src/constants/fileUpload.ts b/frontend/src/constants/fileUpload.ts index d68b2aa3..ae5e60ba 100644 --- a/frontend/src/constants/fileUpload.ts +++ b/frontend/src/constants/fileUpload.ts @@ -160,27 +160,82 @@ const TEXT_SNIFF_BYTES = 8192; const MAX_NONTEXT_RATIO = 0.1; // Control bytes that occur in ordinary text: tab, LF, VT, FF, CR, ESC. const TEXT_CONTROL_BYTES = new Set([0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x1b]); -// A UTF-16/32 file is half NUL bytes, so it has to be recognised by its BOM -// before the NUL test — Notepad's "Unicode" .txt is ordinary text. -const TEXT_BOMS = [ - [0xef, 0xbb, 0xbf], - [0xff, 0xfe], - [0xfe, 0xff], - [0x00, 0x00, 0xfe, 0xff], +const TEXT_CONTROL_CHARS = new Set(['\t', '\n', '\v', '\f', '\r', '\u001b']); +// \p{C}: control, format, private-use, surrogate and unassigned code points — +// what binary decodes into. +const NONTEXT_CHAR = /\p{C}/u; +// A UTF-16/32 file is half NUL bytes, so the byte rules below would reject it. +// Its BOM says which encoding to read it as; the decoded characters are then +// judged instead. Longest BOM first — UTF-32-LE starts with the UTF-16-LE one. +// A BOM is never a verdict on its own: three bytes must not buy a video a pass. +const TEXT_BOMS: { bom: number[]; encoding: string | null }[] = [ + // TextDecoder has no UTF-32; leave those two to the server, which does. + { bom: [0xff, 0xfe, 0x00, 0x00], encoding: null }, + { bom: [0x00, 0x00, 0xfe, 0xff], encoding: null }, + { bom: [0xef, 0xbb, 0xbf], encoding: 'utf-8' }, + { bom: [0xff, 0xfe], encoding: 'utf-16le' }, + { bom: [0xfe, 0xff], encoding: 'utf-16be' }, ]; -/** Whether a leading byte sample reads as text rather than binary. */ +/** + * Whether decoded characters read as text rather than decoded binary. + * Replacement characters count against it, so bytes the decoder could not + * read are evidence rather than silently dropped. + */ +function decodedLooksLikeText(text: string): boolean { + // The byte-level NUL rule, one level up: no text holds a NUL character. + if (text.includes('\u0000')) return false; + let total = 0; + let nontext = 0; + for (const char of text) { + total += 1; + if ( + char === '\ufffd' || + (!TEXT_CONTROL_CHARS.has(char) && NONTEXT_CHAR.test(char)) + ) + nontext += 1; + } + if (total === 0) return true; + return nontext / total <= MAX_NONTEXT_RATIO; +} + +/** + * Whether a leading byte sample reads as text rather than binary. A UTF-16 + * BOM switches the test to the decoded characters; the content is still what + * decides. Mirrors `looks_like_text` in application/upload_limits.py. + */ export function looksLikeText(sample: Uint8Array): boolean { if (sample.length === 0) return true; - if (TEXT_BOMS.some((bom) => bom.every((byte, i) => sample[i] === byte))) - return true; + + let body = sample; + const marked = TEXT_BOMS.find(({ bom }) => + bom.every((byte, i) => sample[i] === byte), + ); + if (marked) { + // UTF-32, or a decoder this browser lacks: defer to the server rather + // than refuse a file it would accept. + if (marked.encoding === null) return true; + body = sample.subarray(marked.bom.length); + if (body.length === 0) return true; + if (marked.encoding !== 'utf-8') { + try { + return decodedLooksLikeText( + new TextDecoder(marked.encoding).decode(body), + ); + } catch { + return true; + } + } + // UTF-8 BOM: the byte rules still apply to everything after it. + } + let nontext = 0; - for (const byte of sample) { + for (const byte of body) { if (byte === 0x00) return false; if ((byte < 0x20 && !TEXT_CONTROL_BYTES.has(byte)) || byte === 0x7f) nontext += 1; } - return nontext / sample.length <= MAX_NONTEXT_RATIO; + return nontext / body.length <= MAX_NONTEXT_RATIO; } async function isSupportedAttachmentFile(file: File): Promise { diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index ab00a543..de72eb93 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1583,5 +1583,26 @@ def test_record_responses_metadata_captures_cache_write_bin(monkeypatch): ), )) # A pure cache-write turn (first request of a prefix) has reads=0 and - # writes>0; the details must still be recorded. - assert llm._last_usage["prompt_tokens_details"] == {"cache_write_tokens": 9} + # writes>0; both bins were reported, so both are recorded — a reported + # zero is "no cache hits", not "unknown". + assert llm._last_usage["prompt_tokens_details"] == { + "cached_tokens": 0, + "cache_write_tokens": 9, + } + + +@pytest.mark.unit +def test_record_responses_metadata_omits_unreported_cache_bins(monkeypatch): + """A provider that says nothing about caching must persist as NULL, not 0.""" + llm = _make_llm(monkeypatch, _responses_caps()) + llm._record_responses_metadata(_ns( + id="resp_2", + usage=_ns( + input_tokens=10, + output_tokens=7, + total_tokens=17, + input_tokens_details=None, + output_tokens_details=None, + ), + )) + assert "prompt_tokens_details" not in llm._last_usage diff --git a/tests/test_upload_limits.py b/tests/test_upload_limits.py index 87c868f3..9d41e45a 100644 --- a/tests/test_upload_limits.py +++ b/tests/test_upload_limits.py @@ -121,6 +121,49 @@ def test_enforce_parseable_attachment_accepts_bom_marked_unicode_text( enforce_parseable_attachment(path, "notes.txt") +@pytest.mark.parametrize( + "bom", + [codecs.BOM_UTF8, codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE, codecs.BOM_UTF32_BE], +) +def test_enforce_parseable_attachment_rejects_binary_behind_a_bom(bom, tmp_path): + """A BOM says which encoding to read, not that the content is text. + + Otherwise three prepended bytes buy any binary a pass. + """ + from application.upload_limits import ( + enforce_parseable_attachment, + UnsupportedUploadTypeError, + ) + + path = tmp_path / "notes.txt" + path.write_bytes(bom + MP4_HEADER + bytes(range(256)) * 8) + + with pytest.raises(UnsupportedUploadTypeError): + enforce_parseable_attachment(path, "notes.txt") + + +def test_enforce_parseable_attachment_uses_the_extractor_it_is_given(tmp_path): + """The worker holds the live parser table; a trimmed install must not admit on trust. + + Without docling the fallback extractor has no .webp handler, so a .webp + would otherwise skip the content check and be read as plain text. + """ + from application.upload_limits import ( + enforce_parseable_attachment, + UnsupportedUploadTypeError, + ) + + path = tmp_path / "scan.webp" + path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256))) + + # Default list: .webp is parser-backed, admitted on its name. + enforce_parseable_attachment(path, "scan.webp") + + # The extractor actually loaded has no .webp parser. + with pytest.raises(UnsupportedUploadTypeError): + enforce_parseable_attachment(path, "scan.webp", {".pdf", ".docx"}) + + @pytest.mark.parametrize( "filename", [ @@ -177,6 +220,10 @@ def test_enforce_parseable_attachment_accepts_text_without_a_parser(filename, tm (b"\x1b[31mred log line\x1b[0m\n", True), (codecs.BOM_UTF16_LE + "hi\n".encode("utf-16-le"), True), (codecs.BOM_UTF8 + b"hi\n", True), + (codecs.BOM_UTF32_LE + "hi\n".encode("utf-32-le"), True), + # A BOM in front of binary is still binary. + (codecs.BOM_UTF8 + b"\x00\x01\x02", False), + (codecs.BOM_UTF16_LE + MP4_HEADER, False), (b"text\x00with nul", False), (bytes(range(32)) * 4, False), (b"\x7f\x7f\x7f\x7f" + b"a" * 16, False), diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index e5790cae..a48ef4ec 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -379,6 +379,53 @@ class TestAttachmentTypeGuard: failed = [payload for name, payload in events if name == "attachment.failed"] assert failed and failed[0]["error"] == "Unsupported file type: .mp4" + def test_suffix_the_loaded_extractor_cannot_parse_is_refused( + self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path + ): + """The guard judges against the parser table actually loaded. + + Without docling the fallback extractor has no .webp parser, so a + .webp the route admitted on its name would otherwise be opened as + plain text here. + """ + from application import worker + + local_path = tmp_path / "scan.webp" + local_path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPVP8 " + bytes(range(256))) + + events = [] + fake_storage = MagicMock(name="storage") + fake_storage.process_file.side_effect = lambda path, callback: callback( + str(local_path) + ) + monkeypatch.setattr(worker.StorageCreator, "get_storage", lambda: fake_storage) + # The docling-less fallback table: images are handled, .webp is not. + monkeypatch.setattr( + worker, + "get_default_file_extractor", + lambda ocr_enabled=False, pdf_text_fast_path=False: {".png": object()}, + ) + monkeypatch.setattr( + worker, + "publish_user_event", + lambda user, name, payload, **kwargs: events.append((name, payload)), + ) + + file_info = { + "filename": "scan.webp", + "attachment_id": "507f1f77bcf86cd799439014", + "path": "uploads/user1/attachments/scan.webp", + "metadata": {"source": "chat"}, + } + + with pytest.raises( + worker.AttachmentRejectedError, match=r"Unsupported file type: \.webp" + ): + worker.attachment_worker(task_self, file_info, "user1") + + failed = [payload for name, payload in events if name == "attachment.failed"] + assert failed and failed[0]["error"] == "Unsupported file type: .webp" + def test_text_without_a_parser_is_parsed( self, pg_conn, patch_worker_db, task_self, monkeypatch, tmp_path ):