"""Attachments: upload validation, extraction, and how they reach the model."""
from __future__ import annotations
import io
from datetime import timedelta
import pytest
from fastapi.testclient import TestClient
from PIL import Image
from sqlalchemy import select
from lembas.db.models import Attachment, Chat, Connection, Message, Model
from lembas.services import chat as chat_service
from lembas.services import files as files_service
from lembas.services import settings_store
from lembas.services.crypto import encrypt
# --- Fixtures ----------------------------------------------------------------
def png_bytes(width: int = 40, height: int = 30, mode: str = "RGB") -> bytes:
buffer = io.BytesIO()
Image.new(mode, (width, height), "red").save(buffer, format="PNG")
return buffer.getvalue()
def jpeg_bytes(width: int = 40, height: int = 30) -> bytes:
buffer = io.BytesIO()
Image.new("RGB", (width, height), "blue").save(buffer, format="JPEG")
return buffer.getvalue()
def pdf_bytes(pages: list[str]) -> bytes:
"""A real PDF with a text layer, built with pypdf + reportlab-free drawing.
pypdf cannot author text, so the file is assembled by hand. It is minimal
but genuinely parseable, which is the point -- a fake would not exercise
extraction at all.
"""
objects: list[bytes] = []
def stream_for(text: str) -> bytes:
escaped = text.replace("\\", r"\\").replace("(", r"\(").replace(")", r"\)")
return f"BT /F1 12 Tf 72 720 Td ({escaped}) Tj ET".encode("latin-1")
page_ids = [4 + i * 2 for i in range(len(pages))]
kids = " ".join(f"{pid} 0 R" for pid in page_ids)
objects.append(b"<< /Type /Catalog /Pages 2 0 R >>")
objects.append(f"<< /Type /Pages /Kids [{kids}] /Count {len(pages)} >>".encode())
objects.append(b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>")
for index, text in enumerate(pages):
content = stream_for(text)
objects.append(
f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
f"/Resources << /Font << /F1 3 0 R >> >> "
f"/Contents {page_ids[index] + 1} 0 R >>".encode()
)
objects.append(
b"<< /Length "
+ str(len(content)).encode()
+ b" >>\nstream\n"
+ content
+ b"\nendstream"
)
out = bytearray(b"%PDF-1.4\n")
offsets = [0]
for number, body in enumerate(objects, start=1):
offsets.append(len(out))
out += f"{number} 0 obj\n".encode() + body + b"\nendobj\n"
xref_at = len(out)
out += f"xref\n0 {len(objects) + 1}\n".encode()
out += b"0000000000 65535 f \n"
for offset in offsets[1:]:
out += f"{offset:010d} 00000 n \n".encode()
out += (
f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\nstartxref\n{xref_at}\n%%EOF".encode()
)
return bytes(out)
@pytest.fixture
def chat_with_model(client: TestClient, db, registered):
"""A chat whose model has vision turned on."""
connection = Connection(name="T", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt(""))
db.add(connection)
db.commit()
db.add(
Model(
connection_id=connection.id,
model_id="seeing-model",
capabilities_json={"vision": True},
)
)
db.commit()
chat_id = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
return chat_id
# --- Type detection and processing -------------------------------------------
def test_png_is_recognised_as_an_image():
prepared = files_service.prepare(png_bytes(), "photo.png")
assert prepared.kind == "image"
assert prepared.width == 40 and prepared.height == 30
def test_large_images_are_downscaled():
"""A phone photo is megabytes of base64 and a large slice of the context."""
prepared = files_service.prepare(jpeg_bytes(4000, 3000), "big.jpg")
assert max(prepared.width, prepared.height) == files_service.MAX_IMAGE_EDGE
assert prepared.height == int(3000 * (files_service.MAX_IMAGE_EDGE / 4000))
def test_small_images_are_left_alone():
prepared = files_service.prepare(jpeg_bytes(100, 80), "small.jpg")
assert (prepared.width, prepared.height) == (100, 80)
def test_transparent_images_stay_png():
prepared = files_service.prepare(png_bytes(mode="RGBA"), "logo.png")
assert prepared.media_type in ("image/png", "image/jpeg")
def test_a_file_lying_about_its_type_is_judged_by_its_bytes():
"""The extension says PNG; the content is text, and content wins."""
prepared = files_service.prepare(b"just some words", "trick.png")
assert prepared.kind == "text"
def test_pdf_text_is_extracted():
prepared = files_service.prepare(pdf_bytes(["Lembas keeps a traveller going."]), "doc.pdf")
assert prepared.kind == "document"
assert prepared.pages == 1
assert "Lembas keeps a traveller going." in prepared.extracted_text
def test_multi_page_pdfs_are_labelled_by_page():
prepared = files_service.prepare(pdf_bytes(["First page here", "Second page here"]), "d.pdf")
assert prepared.pages == 2
assert "[page 1]" in prepared.extracted_text
assert "[page 2]" in prepared.extracted_text
def test_a_pdf_with_no_text_layer_says_so():
"""A scan otherwise looks like the model simply ignored the document."""
prepared = files_service.prepare(pdf_bytes([" "]), "scan.pdf")
assert prepared.extraction_error
assert "scanned" in prepared.extraction_error.lower()
def test_a_corrupt_pdf_is_stored_with_an_error_not_rejected():
prepared = files_service.prepare(b"%PDF-1.4\nthis is not really a pdf", "broken.pdf")
assert prepared.kind == "document"
assert prepared.extraction_error
def test_text_files_are_decoded():
prepared = files_service.prepare(b"line one\nline two", "notes.txt")
assert prepared.kind == "text"
assert prepared.extracted_text == "line one\nline two"
def test_source_files_keep_a_sensible_media_type():
assert files_service.prepare(b"print('hi')", "x.py").media_type == "text/x-python"
def test_binary_files_are_rejected():
with pytest.raises(files_service.FileError):
files_service.prepare(b"\x00\x01\x02\x03" * 100, "mystery.bin")
def test_empty_files_are_rejected():
with pytest.raises(files_service.FileError):
files_service.prepare(b"", "empty.txt")
def test_oversized_files_are_rejected():
with pytest.raises(files_service.FileError) as caught:
files_service.prepare(b"x" * (files_service.MAX_UPLOAD_BYTES + 1), "huge.txt")
assert "MB" in str(caught.value)
def test_extracted_text_is_capped(monkeypatch):
monkeypatch.setattr(files_service, "MAX_EXTRACTED_CHARS", 50)
prepared = files_service.prepare(b"x" * 500, "long.txt")
assert len(prepared.extracted_text) == 50
assert prepared.truncated is True
# --- Path safety -------------------------------------------------------------
@pytest.mark.parametrize(
"name", ["../../etc/passwd", "..\\windows", "/etc/passwd", ".hidden", ""]
)
def test_stored_path_refuses_traversal(name):
assert files_service.stored_path(name) is None
def test_stored_names_are_random_not_the_uploaders(client: TestClient, db, registered):
response = client.post(
"/api/files", files={"file": ("../../evil.txt", b"content", "text/plain")}
)
assert response.status_code == 200
attachment = db.scalar(select(Attachment))
assert "/" not in attachment.stored_name
assert ".." not in attachment.stored_name
# The display name is kept, but only as a label.
assert attachment.filename == "evil.txt"
# --- Upload through the API --------------------------------------------------
def test_upload_returns_a_chip(client: TestClient, db, registered):
response = client.post(
"/api/files", files={"file": ("notes.txt", b"hello there", "text/plain")}
)
assert response.status_code == 200
assert "notes.txt" in response.text
assert 'name="file_ids"' in response.text
def test_a_rejected_upload_returns_a_readable_error(client: TestClient, db, registered):
response = client.post(
"/api/files", files={"file": ("bad.bin", b"\x00\x01" * 500, "application/octet-stream")}
)
assert response.status_code == 200
assert "chip--error" in response.text
assert db.scalar(select(Attachment)) is None
def test_uploading_needs_permission(client: TestClient, db, registered):
client.post("/auth/logout", follow_redirects=False)
client.post(
"/auth/register",
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
follow_redirects=False,
)
settings_store.update(db, {"default_permissions": {"files.upload": False}})
response = client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
assert response.status_code == 403
def test_you_cannot_fetch_someone_elses_file(client: TestClient, db, registered):
client.post("/api/files", files={"file": ("secret.txt", b"mine", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post("/auth/logout", follow_redirects=False)
client.post(
"/auth/register",
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
follow_redirects=False,
)
assert client.get(f"/api/files/{attachment.id}/content").status_code == 404
def test_non_images_are_served_as_downloads(client: TestClient, db, registered):
"""An uploaded .html served inline would run in this origin."""
client.post("/api/files", files={"file": ("page.html", b"hi", "text/html")})
attachment = db.scalar(select(Attachment))
response = client.get(f"/api/files/{attachment.id}/content")
assert "attachment;" in response.headers["content-disposition"]
assert response.headers["content-type"] == "application/octet-stream"
assert response.headers["x-content-type-options"] == "nosniff"
def test_images_are_served_inline(client: TestClient, db, registered):
client.post("/api/files", files={"file": ("p.png", png_bytes(), "image/png")})
attachment = db.scalar(select(Attachment))
response = client.get(f"/api/files/{attachment.id}/content")
assert "inline;" in response.headers["content-disposition"]
def test_an_unsent_attachment_can_be_removed(client: TestClient, db, registered):
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
stored = files_service.stored_path(attachment.stored_name)
assert client.delete(f"/api/files/{attachment.id}").status_code == 200
assert db.scalar(select(Attachment)) is None
assert not stored.exists()
def test_a_sent_attachment_cannot_be_removed(client: TestClient, db, chat_with_model):
"""Deleting it would rewrite a conversation the user has already read."""
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "look", "file_ids": [attachment.id]},
)
assert client.delete(f"/api/files/{attachment.id}").status_code == 409
# --- Attaching to a message --------------------------------------------------
def test_sending_binds_the_attachment_to_the_message(client: TestClient, db, chat_with_model):
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "have a look", "file_ids": [attachment.id]},
)
db.refresh(attachment)
message = db.scalar(select(Message).where(Message.role == "user"))
assert attachment.message_id == message.id
def test_a_message_with_only_an_attachment_is_accepted(client: TestClient, db, chat_with_model):
""""Here, look at this" with no words is a legitimate turn."""
client.post("/api/files", files={"file": ("a.png", png_bytes(), "image/png")})
attachment = db.scalar(select(Attachment))
response = client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "", "file_ids": [attachment.id]},
)
assert response.status_code == 200
assert db.scalar(select(Message).where(Message.role == "user")) is not None
def test_a_truly_empty_message_is_still_ignored(client: TestClient, db, chat_with_model):
response = client.post(f"/api/chats/{chat_with_model}/messages", data={"content": " "})
assert response.status_code == 204
assert db.scalar(select(Message)) is None
def test_you_cannot_attach_someone_elses_file(client: TestClient, db, chat_with_model):
"""A forged id must not pull another user's file into a conversation."""
client.post("/api/files", files={"file": ("mine.txt", b"secret", "text/plain")})
stolen = db.scalar(select(Attachment))
client.post("/auth/logout", follow_redirects=False)
client.post(
"/auth/register",
data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
follow_redirects=False,
)
their_chat = client.post("/api/chats").headers["HX-Redirect"].rsplit("/", 1)[1]
client.post(
f"/api/chats/{their_chat}/messages",
data={"content": "gimme", "file_ids": [stolen.id]},
)
db.refresh(stolen)
assert stolen.message_id is None
# --- What reaches the model --------------------------------------------------
def _user_message(db, chat_id: str) -> Message:
return db.scalar(
select(Message).where(Message.chat_id == chat_id, Message.role == "user")
)
def test_images_become_multimodal_parts_for_a_vision_model(client: TestClient, db, chat_with_model):
client.post("/api/files", files={"file": ("p.png", png_bytes(), "image/png")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "what is this?", "file_ids": [attachment.id]},
)
chat = db.get(Chat, chat_with_model)
payload = chat_service.build_request(db, chat)
content = payload["messages"][0]["content"]
assert isinstance(content, list)
assert content[0] == {"type": "text", "text": "what is this?"}
assert content[1]["type"] == "image_url"
assert content[1]["image_url"]["url"].startswith("data:image/")
def test_images_are_withheld_from_a_model_without_vision(client: TestClient, db, chat_with_model):
"""Most endpoints reject the whole request rather than ignoring the image."""
model = db.scalar(select(Model))
model.capabilities_json = {"vision": False}
db.commit()
client.post("/api/files", files={"file": ("p.png", png_bytes(), "image/png")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "what is this?", "file_ids": [attachment.id]},
)
chat = db.get(Chat, chat_with_model)
content = chat_service.build_request(db, chat)["messages"][0]["content"]
assert isinstance(content, str)
assert content == "what is this?"
def test_a_plain_turn_stays_a_plain_string(client: TestClient, db, chat_with_model):
"""The list form is a reliable 400 from endpoints that do not implement it."""
client.post(f"/api/chats/{chat_with_model}/messages", data={"content": "just words"})
chat = db.get(Chat, chat_with_model)
assert chat_service.build_request(db, chat)["messages"][0]["content"] == "just words"
def test_document_text_is_prepended_in_tags(client: TestClient, db, chat_with_model):
client.post(
"/api/files", files={"file": ("report.txt", b"Quarterly results were good.", "text/plain")}
)
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "summarise this", "file_ids": [attachment.id]},
)
chat = db.get(Chat, chat_with_model)
content = chat_service.build_request(db, chat)["messages"][0]["content"]
assert '' in content
assert "Quarterly results were good." in content
# The question comes after the material it refers to.
assert content.index("") < content.index("summarise this")
def test_documents_reach_a_model_without_vision(client: TestClient, db, chat_with_model):
model = db.scalar(select(Model))
model.capabilities_json = {}
db.commit()
client.post("/api/files", files={"file": ("notes.txt", b"important detail", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "read it", "file_ids": [attachment.id]},
)
chat = db.get(Chat, chat_with_model)
content = chat_service.build_request(db, chat)["messages"][0]["content"]
assert "important detail" in content
def test_truncation_is_declared_to_the_model(client: TestClient, db, chat_with_model, monkeypatch):
"""A model asked about page 400 should be able to say it did not see it."""
monkeypatch.setattr(files_service, "MAX_EXTRACTED_CHARS", 20)
client.post("/api/files", files={"file": ("big.txt", b"y" * 200, "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "read", "file_ids": [attachment.id]},
)
chat = db.get(Chat, chat_with_model)
assert "(truncated)" in chat_service.build_request(db, chat)["messages"][0]["content"]
def test_an_attachment_only_turn_still_reaches_the_model(client: TestClient, db, chat_with_model):
client.post("/api/files", files={"file": ("p.png", png_bytes(), "image/png")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages", data={"content": "", "file_ids": [attachment.id]}
)
chat = db.get(Chat, chat_with_model)
messages = chat_service.build_request(db, chat)["messages"]
assert len(messages) == 1
assert messages[0]["content"][0]["type"] == "image_url"
# --- Housekeeping ------------------------------------------------------------
def test_orphans_are_swept(client: TestClient, db, registered):
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
path = files_service.stored_path(attachment.stored_name)
assert files_service.sweep_orphans(db, timedelta(hours=24)) == 0 # too new
assert files_service.sweep_orphans(db, timedelta(seconds=-1)) == 1
assert db.scalar(select(Attachment)) is None
assert not path.exists()
def test_sent_attachments_are_never_swept(client: TestClient, db, chat_with_model):
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "keep", "file_ids": [attachment.id]},
)
assert files_service.sweep_orphans(db, timedelta(seconds=-1)) == 0
assert db.scalar(select(Attachment)) is not None
def test_deleting_a_message_deletes_its_attachments(client: TestClient, db, chat_with_model):
client.post("/api/files", files={"file": ("a.txt", b"hi", "text/plain")})
attachment = db.scalar(select(Attachment))
client.post(
f"/api/chats/{chat_with_model}/messages",
data={"content": "x", "file_ids": [attachment.id]},
)
client.delete(f"/api/chats/{chat_with_model}")
assert db.scalar(select(Attachment)) is None