odysseus/tests/test_replace_messages_upload_reservations.py

"""Upload lifecycle guarantees for compaction's replace_messages path."""

import concurrent.futures
import json
import os
import threading
import uuid

import pytest
from sqlalchemy import event

import core.database as cdb
import core.session_manager as session_manager_module
from core.models import ChatMessage
from src.upload_handler import UploadHandler
from tests.helpers.sqlite_db import make_temp_sqlite


OLD_TIMESTAMP = "2000-01-01T00:00:00"


@pytest.fixture
def manager_db(monkeypatch):
    SessionLocal, engine, tmpfile = make_temp_sqlite(cdb.Base.metadata)
    monkeypatch.setattr(session_manager_module, "SessionLocal", SessionLocal)
    manager = session_manager_module.SessionManager.__new__(
        session_manager_module.SessionManager
    )
    manager.sessions = {}
    manager.upload_handler = None
    try:
        yield manager, SessionLocal, engine
    finally:
        engine.dispose()
        tmpfile.close()
        try:
            os.unlink(tmpfile.name)
        except OSError:
            pass


def _seed_session(SessionLocal, *, owner="alice", content="existing durable history"):
    session_id = "replace-" + uuid.uuid4().hex
    db = SessionLocal()
    try:
        db.add(cdb.Session(
            id=session_id,
            owner=owner,
            name="Compaction reservation regression",
            model="test-model",
            endpoint_url="http://localhost:11434",
            archived=False,
            message_count=1,
        ))
        db.add(cdb.ChatMessage(
            id="message-" + uuid.uuid4().hex,
            session_id=session_id,
            role="user",
            content=content,
            meta_data=json.dumps({"source": "before-replacement"}),
        ))
        db.commit()
    finally:
        db.close()
    return session_id


def _attachment_message(upload_id, text):
    return ChatMessage(
        role="user",
        content=text,
        metadata={
            "attachments": [{
                "id": upload_id,
                "name": f"{text}.txt",
                "mime": "text/plain",
                "size": len(text),
            }]
        },
    )


def _durable_messages(SessionLocal, session_id):
    db = SessionLocal()
    try:
        return [
            (message.role, message.content, message.meta_data)
            for message in db.query(cdb.ChatMessage)
            .filter(cdb.ChatMessage.session_id == session_id)
            .order_by(cdb.ChatMessage.timestamp, cdb.ChatMessage.id)
            .all()
        ]
    finally:
        db.close()


def test_replace_messages_reserves_every_incoming_attachment_before_delete(
    manager_db,
    monkeypatch,
):
    manager, SessionLocal, engine = manager_db
    session_id = _seed_session(SessionLocal)
    manager.upload_handler = object()
    incoming = [
        _attachment_message("1" * 32 + ".txt", "first"),
        _attachment_message("2" * 32 + ".txt", "second"),
    ]
    message_mutations = []
    reservation_calls = []

    def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany):
        normalized = statement.lstrip().upper()
        if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")):
            message_mutations.append(normalized.split(maxsplit=1)[0])

    def reserve(handler, owner, content, metadata):
        assert message_mutations == []
        reservation_calls.append((handler, owner, content, metadata))
        return None

    event.listen(engine, "before_cursor_execute", record_sql)
    monkeypatch.setattr(
        session_manager_module,
        "reserve_message_upload_references",
        reserve,
    )
    try:
        assert manager.replace_messages(session_id, incoming) is True
    finally:
        event.remove(engine, "before_cursor_execute", record_sql)

    assert [call[2] for call in reservation_calls] == ["first", "second"]
    assert all(call[0] is manager.upload_handler for call in reservation_calls)
    assert all(call[1] == "alice" for call in reservation_calls)
    assert message_mutations == ["DELETE", "INSERT"]


def test_replace_messages_reservation_failure_leaves_durable_history_unchanged(
    manager_db,
    monkeypatch,
):
    manager, SessionLocal, engine = manager_db
    session_id = _seed_session(SessionLocal)
    before = _durable_messages(SessionLocal, session_id)
    manager.upload_handler = object()
    missing_upload_id = "4" * 32 + ".txt"
    incoming = [
        _attachment_message("3" * 32 + ".txt", "available"),
        _attachment_message(missing_upload_id, "missing"),
    ]
    calls = []
    message_mutations = []

    def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany):
        normalized = statement.lstrip().upper()
        if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")):
            message_mutations.append(normalized.split(maxsplit=1)[0])

    def reserve(_handler, _owner, content, _metadata):
        calls.append(content)
        return missing_upload_id if content == "missing" else None

    event.listen(engine, "before_cursor_execute", record_sql)
    monkeypatch.setattr(
        session_manager_module,
        "reserve_message_upload_references",
        reserve,
    )
    try:
        assert manager.replace_messages(session_id, incoming) is False
    finally:
        event.remove(engine, "before_cursor_execute", record_sql)

    assert calls == ["available", "missing"]
    assert message_mutations == []
    assert _durable_messages(SessionLocal, session_id) == before
    assert [message.content for message in manager.sessions[session_id].history] == [
        "existing durable history"
    ]
    assert all("_db_id" not in (message.metadata or {}) for message in incoming)


def test_cleanup_cannot_delete_attachment_during_concurrent_compaction_replacement(
    manager_db,
    monkeypatch,
    tmp_path,
):
    manager, SessionLocal, _engine = manager_db
    session_id = _seed_session(SessionLocal)
    base_dir = tmp_path / "base"
    upload_dir = tmp_path / "uploads"
    base_dir.mkdir()
    upload_dir.mkdir()
    handler = UploadHandler(str(base_dir), str(upload_dir))
    manager.upload_handler = handler

    upload_id = "5" * 32 + ".txt"
    upload_hash = "6" * 64
    dated_dir = upload_dir / "2000" / "01" / "01"
    dated_dir.mkdir(parents=True)
    upload_path = dated_dir / upload_id
    upload_path.write_text("attachment retained by compaction", encoding="utf-8")
    upload_row = {
        "id": upload_id,
        "path": str(upload_path),
        "mime": "text/plain",
        "size": upload_path.stat().st_size,
        "name": "compaction.txt",
        "original_name": "compaction.txt",
        "hash": upload_hash,
        "checksum_sha256": upload_hash,
        "uploaded_at": OLD_TIMESTAMP,
        "created_at": OLD_TIMESTAMP,
        "last_accessed": OLD_TIMESTAMP,
        "owner": "alice",
    }
    (upload_dir / "uploads.json").write_text(
        json.dumps({f"alice:{upload_hash}": upload_row}),
        encoding="utf-8",
    )
    handler._index_cache = None

    reservation_write_entered = threading.Event()
    release_reservation_write = threading.Event()
    real_atomic_write = handler._atomic_write_json

    def block_reservation_write(path, data, *, sync_backup=False):
        refreshed = any(
            isinstance(row, dict)
            and row.get("id") == upload_id
            and row.get("last_accessed") != OLD_TIMESTAMP
            for row in data.values()
        )
        if sync_backup and refreshed and not reservation_write_entered.is_set():
            reservation_write_entered.set()
            assert release_reservation_write.wait(5)
        return real_atomic_write(path, data, sync_backup=sync_backup)

    monkeypatch.setattr(handler, "_atomic_write_json", block_reservation_write)
    incoming = [_attachment_message(upload_id, "retained after compaction")]

    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        replace_future = pool.submit(manager.replace_messages, session_id, incoming)
        assert reservation_write_entered.wait(5)
        cleanup_future = pool.submit(handler.cleanup_old_uploads, set(), set())
        try:
            with pytest.raises(concurrent.futures.TimeoutError):
                cleanup_future.result(timeout=0.1)
        finally:
            release_reservation_write.set()

        assert replace_future.result(timeout=5) is True
        assert cleanup_future.result(timeout=5) == 0

    assert upload_path.is_file()
    assert handler.resolve_upload(upload_id, owner="alice") is not None
    durable = _durable_messages(SessionLocal, session_id)
    assert len(durable) == 1
    assert json.loads(durable[0][2])["attachments"][0]["id"] == upload_id