mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 10:15:21 +00:00
Merge pull request #2721 from arc53/fix/attachment-type-gate
fix(attachments): refuse unparseable chat attachments
This commit is contained in:
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;")
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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>
|
||||
)}
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)"
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in new issue
Block a user