fix(cleanup): preserve transcripts with stale message counts (#6610)

* fix(cleanup): preserve transcripts with stale message counts

* fix(cleanup): accurately attribute freed space and bound recent session query

* fix(cleanup): add portable fallback for SQLite runtimes lacking RETURNING

---------

Co-authored-by: PR #6166 Reviewer <reviewer@odysseus-review.local>
This commit is contained in:
Robert Vind-Gardoș
2026-10-14 11:21:53 +01:00
committed by GitHub
co-authored by PR #6166 Reviewer
parent 353ea00736
commit 501059805f
2 changed files with 473 additions and 35 deletions
+96 -35
View File
@@ -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)
})
@@ -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