diff --git a/src/cleanup_service.py b/src/cleanup_service.py index ec1503d9c..b19e8c5c7 100644 --- a/src/cleanup_service.py +++ b/src/cleanup_service.py @@ -37,6 +37,41 @@ def _apply_owner_filter(query, DbSession, owner: Optional[str]): return query.filter(DbSession.owner == owner) +def _message_count_subquery(db, DbSession, DbChatMessage): + """Count durable rows instead of trusting the derived session counter.""" + from sqlalchemy import func + return db.query(func.count(DbChatMessage.id)).filter( + DbChatMessage.session_id == DbSession.id + ).correlate(DbSession).scalar_subquery() + + +def _recent_session_ids(db, DbSession, owner: Optional[str]) -> set: + """Fetch the bounded set of recently created session IDs to preserve.""" + recent_q = db.query(DbSession.id).order_by(DbSession.created_at.desc()) + recent_q = _apply_owner_filter(recent_q, DbSession, owner) + return {row[0] for row in recent_q.limit(CleanupConfig.PRESERVE_RECENT_COUNT).all()} + + +def _supports_delete_returning(db) -> bool: + """Check if the current database backend and runtime version support DELETE ... RETURNING.""" + dialect = getattr(db.get_bind(), "dialect", None) + if dialect is None: + return False + if dialect.name == "postgresql": + return True + if dialect.name == "sqlite": + # SQLite added RETURNING support in 3.35.0 (March 2021). + version = getattr(dialect, "server_version_info", None) + if version is not None: + return version >= (3, 35, 0) + try: + import sqlite3 + return sqlite3.sqlite_version_info >= (3, 35, 0) + except Exception: + return False + return getattr(dialect, "delete_returning", False) + + async def archive_inactive_sessions(session_manager, owner: Optional[str] = None) -> int: """ Archive sessions that haven't been accessed in the configured number of days. @@ -96,29 +131,26 @@ async def cleanup_old_sessions(session_manager, owner: Optional[str] = None) -> from src.database import SessionLocal, Session as DbSession, ChatMessage as DbChatMessage db = SessionLocal() try: - recent_q = db.query(DbSession).order_by(DbSession.created_at.desc()) - recent_q = _apply_owner_filter(recent_q, DbSession, owner) - all_sessions = recent_q.all() - recent_session_ids = {session.id for session in all_sessions[:CleanupConfig.PRESERVE_RECENT_COUNT]} + recent_session_ids = _recent_session_ids(db, DbSession, owner) base_query = db.query(DbSession).filter( DbSession.archived == True, DbSession.last_accessed < cutoff_date, DbSession.is_important == False, - DbSession.message_count < CleanupConfig.MIN_MESSAGES_TO_KEEP ) base_query = _apply_owner_filter(base_query, DbSession, owner) - - candidate_sessions = base_query.all() + actual_count = _message_count_subquery(db, DbSession, DbChatMessage) + candidate_sessions = base_query.add_columns(actual_count).all() sessions_to_delete = [] + message_counts = {} preserved_count = 0 - for session in candidate_sessions: + for session, message_count in candidate_sessions: if session.id in recent_session_ids: preserved_count += 1 continue - if session.message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP: + if message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP: preserved_count += 1 continue @@ -128,20 +160,52 @@ async def cleanup_old_sessions(session_manager, owner: Optional[str] = None) -> continue sessions_to_delete.append(session) - - for session in sessions_to_delete: - message_count = db.query(DbChatMessage).filter( - DbChatMessage.session_id == session.id - ).count() - space_freed += message_count * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES + message_counts[session.id] = message_count session_ids = [session.id for session in sessions_to_delete] if session_ids: - db.query(DbSession).filter(DbSession.id.in_(session_ids)).delete(synchronize_session=False) - deleted_count = len(session_ids) + # Recheck the real count in the DELETE statement itself. A message + # arriving after candidate selection must retain the protection. + # Using RETURNING accurately identifies rows deleted by this cleanup, + # avoiding misattributing space freed if another transaction independently + # deleted a candidate row. + from sqlalchemy import delete + stmt = delete(DbSession).filter( + DbSession.id.in_(session_ids), + DbSession.archived == True, + DbSession.last_accessed < cutoff_date, + DbSession.is_important == False, + actual_count < CleanupConfig.MIN_MESSAGES_TO_KEEP, + ) + stmt = _apply_owner_filter(stmt, DbSession, owner) + if _supports_delete_returning(db): + deleted_rows = db.execute(stmt.returning(DbSession.id)).scalars().all() + deleted_ids = set(deleted_rows) + deleted_count = len(deleted_ids) + surviving_ids = {row[0] for row in db.query(DbSession.id).filter( + DbSession.id.in_(session_ids) + ).all()} + unaccounted_ids = set(session_ids) - surviving_ids - deleted_ids + space_freed = sum(message_counts[sid] for sid in deleted_ids) * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES + else: + # Portable fallback for legacy SQLite (< 3.35.0) which lacks RETURNING syntax. + res = db.execute(stmt) + deleted_count = res.rowcount or 0 + surviving_ids = {row[0] for row in db.query(DbSession.id).filter( + DbSession.id.in_(session_ids) + ).all()} + deleted_ids = set(session_ids) - surviving_ids + unaccounted_ids = set() + if deleted_count == 0: + space_freed = 0 + elif len(deleted_ids) <= deleted_count: + space_freed = sum(message_counts[sid] for sid in deleted_ids) * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES + else: + space_freed = sum(message_counts[sid] for sid in sorted(deleted_ids)[:deleted_count]) * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES db.commit() - for session_id in session_ids: + cache_evict_ids = deleted_ids | unaccounted_ids + for session_id in cache_evict_ids: if session_id in session_manager.sessions: del session_manager.sessions[session_id] @@ -176,7 +240,7 @@ async def get_cleanup_preview(owner: Optional[str] = None) -> Dict[str, Any]: estimated_space_freed = 0 preserved_sessions = [] - from src.database import SessionLocal, Session as DbSession + from src.database import SessionLocal, Session as DbSession, ChatMessage as DbChatMessage db = SessionLocal() try: archive_q = db.query(DbSession).filter( @@ -184,49 +248,46 @@ async def get_cleanup_preview(owner: Optional[str] = None) -> Dict[str, Any]: DbSession.archived == False ) archive_q = _apply_owner_filter(archive_q, DbSession, owner) - archive_candidates = archive_q.all() + actual_count = _message_count_subquery(db, DbSession, DbChatMessage) + archive_candidates = archive_q.add_columns(actual_count).all() - for session in archive_candidates: + for session, message_count in archive_candidates: sessions_to_archive.append({ "id": session.id, "name": session.name, "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown", - "message_count": session.message_count + "message_count": message_count }) - recent_q = db.query(DbSession).order_by(DbSession.created_at.desc()) - recent_q = _apply_owner_filter(recent_q, DbSession, owner) - all_sessions = recent_q.all() - recent_session_ids = {session.id for session in all_sessions[:CleanupConfig.PRESERVE_RECENT_COUNT]} + recent_session_ids = _recent_session_ids(db, DbSession, owner) base_query = db.query(DbSession).filter( DbSession.archived == True, DbSession.last_accessed < cutoff_delete, DbSession.is_important == False, - DbSession.message_count < CleanupConfig.MIN_MESSAGES_TO_KEEP ) base_query = _apply_owner_filter(base_query, DbSession, owner) - candidate_sessions = base_query.all() + candidate_sessions = base_query.add_columns(actual_count).all() - for session in candidate_sessions: + for session, message_count in candidate_sessions: if session.id in recent_session_ids: preserved_sessions.append({ "id": session.id, "name": session.name, "reason": f"part of last {CleanupConfig.PRESERVE_RECENT_COUNT} sessions", "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown", - "message_count": session.message_count + "message_count": message_count }) continue - if session.message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP: + if message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP: preserved_sessions.append({ "id": session.id, "name": session.name, "reason": f"has {CleanupConfig.MIN_MESSAGES_TO_KEEP}+ messages", "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown", - "message_count": session.message_count + "message_count": message_count }) continue @@ -238,18 +299,18 @@ async def get_cleanup_preview(owner: Optional[str] = None) -> Dict[str, Any]: "name": session.name, "reason": f"contains keyword: {matching_keywords[0]}", "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown", - "message_count": session.message_count + "message_count": message_count }) continue - session_space = session.message_count * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES + session_space = message_count * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES estimated_space_freed += session_space sessions_to_delete.append({ "id": session.id, "name": session.name, "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown", - "message_count": session.message_count, + "message_count": message_count, "estimated_size_kb": round(session_space / 1024, 2) }) diff --git a/tests/test_cleanup_transcript_preservation.py b/tests/test_cleanup_transcript_preservation.py new file mode 100644 index 000000000..e2c27277e --- /dev/null +++ b/tests/test_cleanup_transcript_preservation.py @@ -0,0 +1,377 @@ +"""Cleanup must use durable transcript rows for its deletion safeguards.""" +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from sqlalchemy import create_engine, event, insert +from sqlalchemy.orm import sessionmaker + +import core.database as database +import src.database as legacy_database +import routes.cleanup.cleanup_routes as routes + + +@pytest.fixture +def cleanup_api(tmp_path, monkeypatch): + engine = create_engine( + f"sqlite:///{(tmp_path / 'cleanup.db').as_posix()}", + connect_args={"check_same_thread": False}, + ) + database.Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine, autoflush=False) + monkeypatch.setattr(legacy_database, "SessionLocal", factory) + monkeypatch.setattr(routes, "get_current_user", lambda request: "alice") + manager = SimpleNamespace(sessions={}) + app = FastAPI() + app.include_router(routes.setup_cleanup_routes(manager)) + with TestClient(app) as client: + yield client, factory, manager + engine.dispose() + + +def seed(factory, manager, stored, actual, *, archived=True): + now = datetime.now(timezone.utc).replace(tzinfo=None) + with factory() as db: + for index in range(10): + db.add(database.Session( + id=f"recent-{index}", name="Recent", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(hours=index), + last_accessed=now, message_count=1, + )) + db.add(database.Session( + id="old-chat", name="Project discussion", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(days=50), + last_accessed=now - timedelta(days=30), archived=archived, + is_important=False, message_count=stored, + )) + db.commit() + db.add_all(database.ChatMessage( + id=f"message-{index}", session_id="old-chat", role="user", + content=f"Saved project discussion {index}", timestamp=now - timedelta(days=30), + ) for index in range(actual)) + db.commit() + manager.sessions["old-chat"] = SimpleNamespace(message_count=actual) + + +@pytest.mark.parametrize("stored", [0, 3]) +def test_cleanup_preserves_long_transcript_with_stale_count(cleanup_api, stored): + client, factory, manager = cleanup_api + seed(factory, manager, stored, 40) + preview = client.get("/api/cleanup/preview").json() + assert not preview["sessions_to_delete"] + protected = next(row for row in preview["preserved_sessions"] if row["id"] == "old-chat") + assert protected["message_count"] == 40 + response = client.post("/api/cleanup") + with factory() as db: + actual = db.query(database.ChatMessage).filter_by(session_id="old-chat").count() + survives = db.get(database.Session, "old-chat") is not None + assert response.json()["deleted_count"] == 0 + assert survives and actual == 40, "Cleanup must count actual rows before hard deletion" + + +def test_cleanup_still_removes_genuinely_short_old_chat(cleanup_api): + client, factory, manager = cleanup_api + seed(factory, manager, 3, 3) + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == 1 + with factory() as db: + assert db.get(database.Session, "old-chat") is None + assert db.query(database.ChatMessage).filter_by(session_id="old-chat").count() == 0 + + +def test_new_archive_is_not_deleted_in_same_cleanup(cleanup_api): + client, factory, manager = cleanup_api + seed(factory, manager, 3, 3, archived=False) + preview = client.get("/api/cleanup/preview").json() + assert not preview["sessions_to_delete"] + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == 0 + with factory() as db: + assert db.get(database.Session, "old-chat").archived + + +@pytest.mark.parametrize("actual,expected_deleted", [(19, 1), (20, 0)]) +def test_cleanup_threshold_uses_actual_count_even_when_metadata_is_high(cleanup_api, actual, expected_deleted): + client, factory, manager = cleanup_api + seed(factory, manager, 100, actual) + preview = client.get("/api/cleanup/preview").json() + section = "sessions_to_delete" if expected_deleted else "preserved_sessions" + row = next(row for row in preview[section] if row["id"] == "old-chat") + assert row["message_count"] == actual + if expected_deleted: + assert row["estimated_size_kb"] == 19 * 512 / 1024 + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == expected_deleted + with factory() as db: + assert (db.get(database.Session, "old-chat") is None) == bool(expected_deleted) + + +@pytest.mark.parametrize("protection", ["starred", "keyword", "other_owner"]) +def test_cleanup_retains_existing_protections(cleanup_api, protection): + client, factory, manager = cleanup_api + seed(factory, manager, 0, 3) + with factory() as db: + row = db.get(database.Session, "old-chat") + if protection == "starred": + row.is_important = True + elif protection == "keyword": + row.name = "Important project" + else: + row.owner = "bob" + # Preserve the old activity date despite the column's onupdate default. + row.last_accessed = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(days=30) + db.commit() + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == 0 + with factory() as db: + assert db.query(database.ChatMessage).filter_by(session_id="old-chat").count() == 3 + + +def test_cleanup_rechecks_messages_at_the_delete_boundary(cleanup_api): + """Deliver new rows after selection but before the DELETE executes.""" + client, factory, manager = cleanup_api + seed(factory, manager, 3, 3) + engine = factory.kw["bind"] + injected = False + + def deliver_messages(connection, cursor, statement, parameters, context, executemany): + nonlocal injected + if not injected and statement.lstrip().startswith("DELETE FROM sessions"): + injected = True + connection.execute(insert(database.ChatMessage), [ + {"id": f"incoming-{i}", "session_id": "old-chat", "role": "assistant", + "content": "New durable message"} for i in range(20) + ]) + + event.listen(engine, "before_cursor_execute", deliver_messages) + try: + response = client.post("/api/cleanup") + finally: + event.remove(engine, "before_cursor_execute", deliver_messages) + assert injected + assert response.json()["deleted_count"] == 0 + assert response.json()["space_freed_mb"] == 0 + assert "old-chat" in manager.sessions + with factory() as db: + assert db.get(database.Session, "old-chat") is not None + assert db.query(database.ChatMessage).filter_by(session_id="old-chat").count() == 23 + + +def test_cleanup_space_freed_ignores_concurrently_deleted_session(cleanup_api): + """If another transaction deletes a candidate first, do not attribute its space.""" + client, factory, manager = cleanup_api + now = datetime.now(timezone.utc).replace(tzinfo=None) + with factory() as db: + for index in range(10): + db.add(database.Session( + id=f"recent-{index}", name="Recent", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(hours=index), + last_accessed=now, message_count=1, + )) + db.add(database.Session( + id="chat-to-delete", name="Eligible 1", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(days=50), + last_accessed=now - timedelta(days=30), archived=True, + is_important=False, message_count=3, + )) + db.add(database.Session( + id="chat-concurrently-deleted", name="Eligible 2", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(days=50), + last_accessed=now - timedelta(days=30), archived=True, + is_important=False, message_count=10, + )) + db.commit() + db.add_all(database.ChatMessage( + id=f"msg1-{i}", session_id="chat-to-delete", role="user", + content="m", timestamp=now - timedelta(days=30), + ) for i in range(3)) + db.add_all(database.ChatMessage( + id=f"msg2-{i}", session_id="chat-concurrently-deleted", role="user", + content="m", timestamp=now - timedelta(days=30), + ) for i in range(10)) + db.commit() + + manager.sessions["chat-to-delete"] = SimpleNamespace(message_count=3) + manager.sessions["chat-concurrently-deleted"] = SimpleNamespace(message_count=10) + + engine = factory.kw["bind"] + injected = False + + def delete_candidate(connection, cursor, statement, parameters, context, executemany): + nonlocal injected + if not injected and statement.lstrip().startswith("DELETE FROM sessions"): + injected = True + with factory() as other_db: + other_db.query(database.Session).filter_by(id="chat-concurrently-deleted").delete() + other_db.commit() + + event.listen(engine, "before_cursor_execute", delete_candidate) + try: + response = client.post("/api/cleanup") + finally: + event.remove(engine, "before_cursor_execute", delete_candidate) + + assert injected + data = response.json() + assert data["deleted_count"] == 1 + # Only chat-to-delete's 3 messages (1536 bytes) freed by this cleanup, not (3 + 10) * 512 + expected_space_mb = round(3 * 512 / (1024 * 1024), 2) + assert data["space_freed_mb"] == expected_space_mb + assert "chat-to-delete" not in manager.sessions + assert "chat-concurrently-deleted" not in manager.sessions + with factory() as db: + assert db.get(database.Session, "chat-to-delete") is None + assert db.get(database.Session, "chat-concurrently-deleted") is None + + +def test_cleanup_recent_session_query_is_bounded(cleanup_api): + """Ensure recent-session preservation scales to large collections and respects limits.""" + client, factory, manager = cleanup_api + now = datetime.now(timezone.utc).replace(tzinfo=None) + with factory() as db: + for index in range(25): + db.add(database.Session( + id=f"sess-{index:02d}", name=f"Session {index}", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(hours=index), + last_accessed=now - timedelta(days=30), archived=True, + is_important=False, message_count=3, + )) + db.add_all(database.ChatMessage( + id=f"msg-{index}-{m}", session_id=f"sess-{index:02d}", role="user", + content="m", timestamp=now - timedelta(days=30), + ) for m in range(3)) + db.commit() + + preview = client.get("/api/cleanup/preview").json() + preserved_ids = {row["id"] for row in preview["preserved_sessions"]} + to_delete_ids = {row["id"] for row in preview["sessions_to_delete"]} + + # Top 10 newest (sess-00 to sess-09) must be preserved + expected_preserved = {f"sess-{i:02d}" for i in range(10)} + expected_to_delete = {f"sess-{i:02d}" for i in range(10, 25)} + assert preserved_ids == expected_preserved + assert to_delete_ids == expected_to_delete + + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == 15 + + with factory() as db: + surviving = {row[0] for row in db.query(database.Session.id).all()} + assert surviving == expected_preserved + + +def test_cleanup_handles_zero_actual_messages_session(cleanup_api): + """Sessions with zero durable messages are properly previewed and cleaned up.""" + client, factory, manager = cleanup_api + seed(factory, manager, 10, 0) + preview = client.get("/api/cleanup/preview").json() + row = next(r for r in preview["sessions_to_delete"] if r["id"] == "old-chat") + assert row["message_count"] == 0 + assert row["estimated_size_kb"] == 0.0 + + response = client.post("/api/cleanup") + assert response.json()["deleted_count"] == 1 + with factory() as db: + assert db.get(database.Session, "old-chat") is None + + +def test_cleanup_db_error_preserves_cache_and_aborts_cleanly(cleanup_api): + """Database execution failures must abort, rollback, and leave in-memory cache intact.""" + client, factory, manager = cleanup_api + seed(factory, manager, 3, 3) + engine = factory.kw["bind"] + injected = False + + def break_delete(connection, cursor, statement, parameters, context, executemany): + nonlocal injected + if not injected and statement.lstrip().startswith("DELETE FROM sessions"): + injected = True + raise RuntimeError("Injected database error during DELETE") + + event.listen(engine, "before_cursor_execute", break_delete) + try: + response = client.post("/api/cleanup") + finally: + event.remove(engine, "before_cursor_execute", break_delete) + + assert injected + assert response.status_code == 200 + assert response.json()["deleted_count"] == 0 + assert response.json()["space_freed_mb"] == 0 + # In-memory session was not prematurely evicted + assert "old-chat" in manager.sessions + # Database row was rolled back and still exists + with factory() as db: + assert db.get(database.Session, "old-chat") is not None + assert db.query(database.ChatMessage).filter_by(session_id="old-chat").count() == 3 + + +def test_cleanup_portable_fallback_without_returning(cleanup_api, monkeypatch): + """Engines lacking RETURNING (e.g. SQLite < 3.35.0) safely clean up without syntax errors.""" + import src.cleanup_service as cleanup_service + monkeypatch.setattr(cleanup_service, "_supports_delete_returning", lambda db: False) + + client, factory, manager = cleanup_api + now = datetime.now(timezone.utc).replace(tzinfo=None) + with factory() as db: + for index in range(10): + db.add(database.Session( + id=f"recent-{index}", name="Recent", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(hours=index), + last_accessed=now, message_count=1, + )) + db.add(database.Session( + id="chat-to-delete", name="Eligible 1", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(days=50), + last_accessed=now - timedelta(days=30), archived=True, + is_important=False, message_count=3, + )) + db.add(database.Session( + id="chat-concurrently-deleted", name="Eligible 2", endpoint_url="", model="", + owner="alice", created_at=now - timedelta(days=50), + last_accessed=now - timedelta(days=30), archived=True, + is_important=False, message_count=10, + )) + db.commit() + db.add_all(database.ChatMessage( + id=f"msg1-{i}", session_id="chat-to-delete", role="user", + content="m", timestamp=now - timedelta(days=30), + ) for i in range(3)) + db.add_all(database.ChatMessage( + id=f"msg2-{i}", session_id="chat-concurrently-deleted", role="user", + content="m", timestamp=now - timedelta(days=30), + ) for i in range(10)) + db.commit() + + manager.sessions["chat-to-delete"] = SimpleNamespace(message_count=3) + manager.sessions["chat-concurrently-deleted"] = SimpleNamespace(message_count=10) + + engine = factory.kw["bind"] + injected = False + + def delete_candidate(connection, cursor, statement, parameters, context, executemany): + nonlocal injected + if not injected and statement.lstrip().startswith("DELETE FROM sessions"): + injected = True + with factory() as other_db: + other_db.query(database.Session).filter_by(id="chat-concurrently-deleted").delete() + other_db.commit() + + event.listen(engine, "before_cursor_execute", delete_candidate) + try: + response = client.post("/api/cleanup") + finally: + event.remove(engine, "before_cursor_execute", delete_candidate) + + assert injected + data = response.json() + assert data["deleted_count"] == 1 + expected_space_mb = round(3 * 512 / (1024 * 1024), 2) + assert data["space_freed_mb"] == expected_space_mb + assert "chat-to-delete" not in manager.sessions + assert "chat-concurrently-deleted" not in manager.sessions + with factory() as db: + assert db.get(database.Session, "chat-to-delete") is None + assert db.get(database.Session, "chat-concurrently-deleted") is None