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/api/user/attachments/routes.py b/application/api/user/attachments/routes.py index f3e7a8b5..3d856cc5 100644 --- a/application/api/user/attachments/routes.py +++ b/application/api/user/attachments/routes.py @@ -39,6 +39,10 @@ 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, + unsupported_upload_message, upload_limit_message, UploadTooLargeError, ) @@ -125,11 +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 unsupported_upload_message(filename) return "Failed to process file" @@ -214,6 +234,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, @@ -242,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: @@ -263,6 +289,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/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 8e1aabc6..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. @@ -710,6 +733,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 +754,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 +783,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 +801,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..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,19 +1305,38 @@ 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) - reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) - except (TypeError, ValueError): - cached = 0 - reasoning = 0 - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + 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 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 | None, written: int | None) -> dict: + """Build the ``prompt_tokens_details`` breakdown from the two cache bins. + + 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 is not None: + details["cached_tokens"] = cached + if written is not None: + details["cache_write_tokens"] = written + return details + @staticmethod def _function_call_ids(response): """call_ids of every ``function_call`` item in a Responses output. @@ -1356,10 +1375,16 @@ 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) - reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + 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 if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result diff --git a/application/parser/file/constants.py b/application/parser/file/constants.py index 4656727b..f8628694 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,69 @@ SUPPORTED_SOURCE_EXTENSIONS = ( *SUPPORTED_SOURCE_IMAGE_EXTENSIONS, *SUPPORTED_AUDIO_EXTENSIONS, ) + +# 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 +# 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. +ATTACHMENT_PARSER_EXTENSIONS = frozenset( + { + *SUPPORTED_SOURCE_EXTENSIONS, + ".xhtml", + ".adoc", + ".asciidoc", + ".tiff", + ".tif", + ".bmp", + ".webp", + ".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"} +) + + +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 ``ATTACHMENT_PARSER_EXTENSIONS``. + """ + return attachment_extension(filename) in ATTACHMENT_PARSER_EXTENSIONS 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/upload_limits.py b/application/upload_limits.py index 98736c29..7ea3b0a8 100644 --- a/application/upload_limits.py +++ b/application/upload_limits.py @@ -2,12 +2,17 @@ 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 ( + attachment_extension, + has_attachment_parser, +) _COPY_CHUNK_BYTES = 64 * 1024 @@ -26,6 +31,169 @@ 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}) +_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 = ( + (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. + + 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. A UTF-16/32 BOM switches the test to + the decoded characters; the content is still what decides. + """ + if not sample: + return True + 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 body: + return False + nontext = sum( + 1 + for byte in body + if (byte < 0x20 and byte not in _TEXT_CONTROL_BYTES) or byte == 0x7F + ) + return nontext / len(body) <= _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, + parser_extensions: Container[str] | None = None, +) -> None: + """Reject an attachment that no parser handles and that is not plain text. + + 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 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 + 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/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/application/worker.py b/application/worker.py index 0e7551f0..168985c6 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,30 @@ class AttachmentRejectedError(Exception): """ +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, 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, parser_extensions) + 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 +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, 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/components/MessageInput.tsx b/frontend/src/components/MessageInput.tsx index 2dc6d364..c472bc58 100644 --- a/frontend/src/components/MessageInput.tsx +++ b/frontend/src/components/MessageInput.tsx @@ -48,7 +48,10 @@ import { useArmedSend } from './message-input/armedSend'; import { handleAbort } from '../conversation/conversationSlice'; import { AUDIO_FILE_ACCEPT_ATTR, - FILE_UPLOAD_ACCEPT, + getFileExtension, + parseUploadErrorMessage, + parseUploadErrorsByIndex, + partitionAttachmentFiles, } from '../constants/fileUpload'; import { UserToolType } from '../settings/types'; import { isChatToolVisible } from '../utils/toolUtils'; @@ -497,8 +500,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 +605,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 +647,10 @@ export default function MessageInput({ dispatch( updateAttachment({ id: uiId, - updates: { status: 'failed' }, + updates: { + status: 'failed', + errorMessage: errorsByIndex.get(index), + }, }), ); return; @@ -748,11 +782,20 @@ export default function MessageInput({ } } else { console.error('Upload failed', status, 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' }, + updates: { + status: 'failed', + errorMessage: + errorsByIndex.get(Number(index)) ?? fallbackMessage, + }, }), ), ); @@ -871,7 +914,10 @@ export default function MessageInput({ dispatch( updateAttachment({ id: uniqueId, - updates: { status: 'failed' }, + updates: { + status: 'failed', + errorMessage: parseUploadErrorMessage(xhr.responseText), + }, }), ); } @@ -891,7 +937,7 @@ export default function MessageInput({ xhr.send(formData); }); }, - [dispatch, token, trackAttachment], + [dispatch, t, token, trackAttachment], ); const handleFileAttachment = (e: React.ChangeEvent) => { @@ -928,7 +974,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 +1643,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..c3fa4839 --- /dev/null +++ b/frontend/src/constants/fileUpload.test.ts @@ -0,0 +1,240 @@ +import { describe, expect, it } from 'vitest'; + +import { + ATTACHMENT_FILE_ACCEPT_ATTR, + ATTACHMENT_PARSER_EXTENSIONS, + getFileExtension, + hasAttachmentParser, + looksLikeText, + parseUploadErrorMessage, + parseUploadErrorsByIndex, + 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); + // .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', () => { + expect(ATTACHMENT_PARSER_EXTENSIONS).toContain('.pdf'); + // Parser-backed suffixes the first cut of this list missed. + for (const ext of [ + '.webp', + '.tiff', + '.tif', + '.bmp', + '.vtt', + '.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); + } + // 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/*'); + }); +}); + +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('accepts BOM-marked Unicode, which is half NUL bytes but still text', () => { + // 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); + 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('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', + ]); + 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), + // 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', 'bom.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([ + 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(); + }); +}); + +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 e8d6f501..ae5e60ba 100644 --- a/frontend/src/constants/fileUpload.ts +++ b/frontend/src/constants/fileUpload.ts @@ -80,3 +80,229 @@ export const SOURCE_FILE_TREE_ACCEPT_ATTR = [ '.jpg', '.jpeg', ].join(','); + +/** + * Chat-attachment suffixes with a dedicated parser. Mirrors the backend's + * `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, 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', + '.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, + * 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 { + 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]); +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 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; + + 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 body) { + if (byte === 0x00) return false; + if ((byte < 0x20 && !TEXT_CONTROL_BYTES.has(byte)) || byte === 0x7f) + nontext += 1; + } + return nontext / body.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; + } +} + +/** + * 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/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..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.""" @@ -2386,3 +2400,120 @@ 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_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 + ): + """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/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..de72eb93 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1550,3 +1550,59 @@ 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; 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/parser/file/test_constants.py b/tests/parser/file/test_constants.py new file mode 100644 index 00000000..f0cf3919 --- /dev/null +++ b/tests/parser/file/test_constants.py @@ -0,0 +1,109 @@ +"""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_SOURCE_EXTENSIONS, +) + +FRONTEND_CONSTANTS = ( + Path(__file__).resolve().parents[3] / "frontend/src/constants/fileUpload.ts" +) + + +@pytest.mark.unit +def test_parser_extensions_match_the_extractor(): + """The list is a hand-kept literal; this is what keeps it honest. + + ``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(ATTACHMENT_PARSER_EXTENSIONS), ( + "parsers exist for suffixes the attachment gate refuses: " + 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_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 + 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") + + 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(ATTACHMENT_PARSER_EXTENSIONS), ( + "frontend and backend attachment lists disagree — " + 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_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 +@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") + # .txt is the plain-text fallthrough itself, not a parser. + assert not has_attachment_parser("notes.txt") 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_upload_limits.py b/tests/test_upload_limits.py index 46818421..9d41e45a 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 @@ -37,3 +38,222 @@ 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") + + +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( + "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", + [ + "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", + [ + "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.""" + 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), + (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), + ], +) +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/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 diff --git a/tests/worker/test_attachment_worker.py b/tests/worker/test_attachment_worker.py index 93e34d97..a48ef4ec 100644 --- a/tests/worker/test_attachment_worker.py +++ b/tests/worker/test_attachment_worker.py @@ -326,3 +326,133 @@ 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_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 + ): + 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