Merge pull request #2721 from arc53/fix/attachment-type-gate

fix(attachments): refuse unparseable chat attachments
This commit is contained in:
Alex authored and GitHub committed 2026-09-03 10:10:45 +01:00
commit a35d41b977
28 files changed
+2049 -125

No files matched your search

@@ -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;")
+44 -2
View File
@@ -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,
+10 -6
View File
@@ -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
+29
View File
@@ -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)
+38 -13
View File
@@ -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
+68
View File
@@ -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
+6
View File
@@ -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(
@@ -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,
+169 -1
View File
@@ -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."""
+28 -1
View File
@@ -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:
+29
View File
@@ -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:
+65 -10
View File
@@ -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<number, string | undefined>();
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<HTMLInputElement>) => {
@@ -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}
/>
<div className="border-border bg-card relative flex w-full flex-col rounded-3xl border dark:bg-transparent">
{/* translate="no": keep Chrome's page translator out of the composer —
it rewrites text nodes into <font> wrappers and React loses the
controls (dead Attach button, see AttachFileButton). */}
<div
translate="no"
className="border-border bg-card relative flex w-full flex-col rounded-3xl border dark:bg-transparent"
>
<AttachmentChipList
attachments={attachments}
draggingId={draggingId}
@@ -1,7 +1,7 @@
import { useTranslation } from 'react-i18next';
import ClipIcon from '../../assets/clip.svg';
import { FILE_UPLOAD_ACCEPT_ATTR } from '../../constants/fileUpload';
import { ATTACHMENT_FILE_ACCEPT_ATTR } from '../../constants/fileUpload';
type AttachFileButtonProps = {
onChange: (e: React.ChangeEvent<HTMLInputElement>) => void;
@@ -11,7 +11,14 @@ export default function AttachFileButton({ onChange }: AttachFileButtonProps) {
const { t } = useTranslation();
return (
<label className="xs:px-3 xs:py-1.5 dark:border-border border-border hover:bg-muted dark:hover:bg-muted flex cursor-pointer items-center rounded-full border px-2 py-1 transition-colors">
// translate="no": Chrome's page translator rewrites text nodes inside
// this label into <font> wrappers and React then loses the control — a
// user with auto-translate on rage-clicked a dead Attach button until
// they reloaded. Keep the composer's controls out of the translator.
<label
translate="no"
className="xs:px-3 xs:py-1.5 dark:border-border border-border hover:bg-muted dark:hover:bg-muted flex cursor-pointer items-center rounded-full border px-2 py-1 transition-colors"
>
<img
src={ClipIcon}
alt="Attach"
@@ -24,7 +31,7 @@ export default function AttachFileButton({ onChange }: AttachFileButtonProps) {
type="file"
className="hidden"
multiple
accept={FILE_UPLOAD_ACCEPT_ATTR}
accept={ATTACHMENT_FILE_ACCEPT_ATTR}
onChange={onChange}
/>
</label>
@@ -25,94 +25,122 @@ export default function AttachmentChipList({
}: AttachmentChipListProps) {
const { t } = useTranslation();
// A tooltip is the one place a touch user can never look, and this list is
// where a phone picker's unsupported file lands. Show the reason inline,
// as soon as it is known, rather than only once a send is attempted.
const failures = attachments.filter(
(attachment) => attachment.status === 'failed' && attachment.errorMessage,
);
return (
<div className="flex flex-wrap gap-1.5 px-2 py-2 sm:gap-2 sm:px-3">
{attachments.map((attachment) => {
return (
<div
key={attachment.id}
draggable={true}
onDragStart={(e) => onDragStart(e, attachment.id)}
onDragOver={onDragOver}
onDrop={(e) => onDropOn(e, attachment.id)}
className={`group dark:text-foreground bg-muted text-muted-foreground dark:bg-accent relative flex items-center rounded-xl px-2 py-1 text-xs sm:px-3 sm:py-1.5 sm:text-sm ${
attachment.status !== 'completed' ? 'opacity-70' : 'opacity-100'
} ${
draggingId === attachment.id
? 'ring-dashed opacity-60 ring-2 ring-purple-200'
: ''
}`}
title={attachment.fileName}
>
<div className="bg-primary mr-2 flex h-8 w-8 items-center justify-center rounded-md p-1">
{attachment.status === 'completed' && (
<img
src={DocumentationDark}
alt="Attachment"
className="h-[15px] w-[15px] object-fill"
/>
)}
{attachment.status === 'failed' && (
<img
src={AlertIcon}
alt="Failed"
className="h-[15px] w-[15px] object-fill"
/>
)}
{(attachment.status === 'uploading' ||
attachment.status === 'processing') && (
<div className="flex h-[15px] w-[15px] items-center justify-center">
<svg className="h-[15px] w-[15px]" viewBox="0 0 24 24">
<circle
className="opacity-0"
cx="12"
cy="12"
r="10"
stroke="transparent"
strokeWidth="4"
fill="none"
/>
<circle
className="text-[#ECECF1]"
cx="12"
cy="12"
r="10"
stroke="currentColor"
strokeWidth="4"
fill="none"
strokeDasharray="62.83"
strokeDashoffset={62.83 * (1 - attachment.progress / 100)}
transform="rotate(-90 12 12)"
/>
</svg>
</div>
)}
</div>
<span className="max-w-[120px] truncate font-medium sm:max-w-[150px]">
{attachment.fileName}
</span>
<Button
type="button"
variant="ghost"
size="icon-sm"
className="ml-1.5 h-auto w-auto rounded-full p-1"
onClick={() => {
onRemove(attachment.id);
}}
aria-label={t('conversation.attachments.remove')}
<>
<div className="flex flex-wrap gap-1.5 px-2 py-2 sm:gap-2 sm:px-3">
{attachments.map((attachment) => {
return (
<div
key={attachment.id}
draggable={true}
onDragStart={(e) => onDragStart(e, attachment.id)}
onDragOver={onDragOver}
onDrop={(e) => onDropOn(e, attachment.id)}
className={`group dark:text-foreground bg-muted text-muted-foreground dark:bg-accent relative flex items-center rounded-xl px-2 py-1 text-xs sm:px-3 sm:py-1.5 sm:text-sm ${
attachment.status !== 'completed' ? 'opacity-70' : 'opacity-100'
} ${
draggingId === attachment.id
? 'ring-dashed opacity-60 ring-2 ring-purple-200'
: ''
}`}
title={
attachment.status === 'failed' && attachment.errorMessage
? `${attachment.fileName}: ${attachment.errorMessage}`
: attachment.fileName
}
>
<X
<div className="bg-primary mr-2 flex h-8 w-8 items-center justify-center rounded-md p-1">
{attachment.status === 'completed' && (
<img
src={DocumentationDark}
alt="Attachment"
className="h-[15px] w-[15px] object-fill"
/>
)}
{attachment.status === 'failed' && (
<img
src={AlertIcon}
alt="Failed"
className="h-[15px] w-[15px] object-fill"
/>
)}
{(attachment.status === 'uploading' ||
attachment.status === 'processing') && (
<div className="flex h-[15px] w-[15px] items-center justify-center">
<svg className="h-[15px] w-[15px]" viewBox="0 0 24 24">
<circle
className="opacity-0"
cx="12"
cy="12"
r="10"
stroke="transparent"
strokeWidth="4"
fill="none"
/>
<circle
className="text-[#ECECF1]"
cx="12"
cy="12"
r="10"
stroke="currentColor"
strokeWidth="4"
fill="none"
strokeDasharray="62.83"
strokeDashoffset={
62.83 * (1 - attachment.progress / 100)
}
transform="rotate(-90 12 12)"
/>
</svg>
</div>
)}
</div>
<span className="max-w-[120px] truncate font-medium sm:max-w-[150px]">
{attachment.fileName}
</span>
<Button
type="button"
variant="ghost"
size="icon-sm"
className="ml-1.5 h-auto w-auto rounded-full p-1"
onClick={() => {
onRemove(attachment.id);
}}
aria-label={t('conversation.attachments.remove')}
className="h-2.5 w-2.5"
/>
</Button>
</div>
);
})}
</div>
>
<X
aria-label={t('conversation.attachments.remove')}
className="h-2.5 w-2.5"
/>
</Button>
</div>
);
})}
</div>
{failures.length > 0 && (
<div
className="flex flex-col gap-0.5 px-2 pb-1 text-xs text-[#B42318] sm:px-3"
role="alert"
>
{failures.map((attachment) => (
<span key={attachment.id}>
{attachment.fileName}: {attachment.errorMessage}
</span>
))}
</div>
)}
</>
);
}
+240
View File
@@ -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('<html>502</html>')).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('<html>502</html>').size).toBe(0);
expect(parseUploadErrorsByIndex('{"message":"nope"}').size).toBe(0);
expect(parseUploadErrorsByIndex('').size).toBe(0);
expect(
parseUploadErrorsByIndex('{"errors":[{"error":"no index"}]}').size,
).toBe(0);
});
});
+226
View File
@@ -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<boolean> {
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<number, string> {
const byIndex = new Map<number, string>();
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;
}
+2 -1
View File
@@ -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",
+2
View File
@@ -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 =
+131
View File
@@ -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
@@ -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:
+35
View File
@@ -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")
+56
View File
@@ -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
+109
View File
@@ -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")
@@ -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)
+100
View File
@@ -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)
+220
View File
@@ -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)"
+112
View File
@@ -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
+130
View File
@@ -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