fix(stabilization): harden attachment lifecycle and agent guard signals (#5420)

* fix: harden stabilization attachment and agent guards * fix(uploads): preserve durable references during cleanup * fix(uploads): close cleanup and compaction races
2026-07-12 12:37:32 +00:00 · 2026-07-11 10:14:14 -04:00
parent 6b8f84553c
commit d16b849c3e
29 changed files with 2718 additions and 189 deletions
@@ -0,0 +1,259 @@
+"""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