From 69319c9390f5829b9e3f0d12a0744f09677ce502 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 14 Jul 2026 15:37:40 +0300
Subject: [PATCH] Grow vacuum retention to protect tagged versions

---
 haiku_rag_slim/haiku/rag/store/engine.py | 47 ++++++++++++++++++++----
 tests/store/test_tags.py                 | 34 +++++++++++++++++
 2 files changed, 73 insertions(+), 8 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py
index bdc18b93..e0986add 100644
--- a/haiku_rag_slim/haiku/rag/store/engine.py
+++ b/haiku_rag_slim/haiku/rag/store/engine.py
@@ -189,6 +189,10 @@ REQUIRED_TABLES: tuple[str, ...] = (
     "settings",
 )
 
+# Keeps the vacuum cleanup cutoff safely older than the oldest tagged
+# version; guards against timestamp precision at the boundary.
+TAG_RETENTION_MARGIN = timedelta(seconds=1)
+
 
 @dataclass
 class TagInfo:
@@ -523,18 +527,45 @@ class Store:
                     retention_seconds = self._config.storage.vacuum_retention_seconds
                 # Perform maintenance per table using optimize() with configurable retention
                 retention = timedelta(seconds=retention_seconds)
-                for table in [
-                    self.documents_table,
-                    self.document_meta_table,
-                    self.chunks_table,
-                    self.document_items_table,
-                    self.settings_table,
-                ]:
-                    await table.optimize(cleanup_older_than=retention)
+                for table in self._tables().values():
+                    await table.optimize(
+                        cleanup_older_than=await self._tag_safe_retention(
+                            table, retention
+                        )
+                    )
             except (RuntimeError, OSError) as e:
                 # Handle resource errors gracefully
                 logger.debug(f"Vacuum skipped due to resource constraints: {e}")
 
+    async def _tag_safe_retention(
+        self, table: lancedb.AsyncTable, retention: timedelta
+    ) -> timedelta:
+        """Grow the retention so the cleanup cutoff stays older than the
+        table's oldest tagged version.
+
+        Lance hard-errors when a tagged version falls inside the cleanup
+        window and the Python API exposes no way to skip tagged versions, so
+        everything older than the oldest tag is retained until that tag is
+        deleted.
+        """
+        tags = await table.tags.list()
+        if not tags:
+            return retention
+
+        timestamps = {v["version"]: v["timestamp"] for v in await table.list_versions()}
+        tagged = [
+            timestamps[tag["version"]]
+            for tag in tags.values()
+            if tag["version"] in timestamps
+        ]
+        if not tagged:
+            return retention
+
+        # LanceDB version timestamps are naive datetimes in local time.
+        oldest = min(ts.replace(tzinfo=None) for ts in tagged)
+        needed = datetime.now() - oldest + TAG_RETENTION_MARGIN
+        return max(retention, needed)
+
     @property
     def _connection_mode(self) -> ConnectionMode:
         return ConnectionMode.from_config(self._config)
diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py
index 04c30e91..e8d010e4 100644
--- a/tests/store/test_tags.py
+++ b/tests/store/test_tags.py
@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 from lancedb.table import AsyncTags
 
@@ -110,6 +112,38 @@ async def test_delete_tag_missing_raises(temp_db_path):
             await store.delete_tag("nope")
 
 
+@pytest.mark.asyncio
+async def test_vacuum_cleans_untagged_versions_and_keeps_tagged(temp_db_path):
+    """Vacuum must both preserve tagged versions (lance hard-errors when a
+    tagged version falls inside the cleanup window, which vacuum would
+    swallow) and still clean untagged versions older than the oldest tag's
+    safety margin."""
+    async with Store(temp_db_path, create=True) as store:
+        repo = DocumentRepository(store)
+        await repo.create(Document(content="First document"))
+        versions_before = [
+            v["version"] for v in await store.list_table_versions("documents")
+        ]
+
+        # Age the pre-tag versions past the retention safety margin.
+        await asyncio.sleep(1.5)
+
+        await repo.create(Document(content="Second document"))
+        await store.create_tag("release-1")
+        tagged_version = (await store.list_tags())["release-1"].tables["documents"]
+
+        await store.vacuum(retention_seconds=0)
+
+        remaining = [v["version"] for v in await store.list_table_versions("documents")]
+        assert tagged_version in remaining
+        assert min(versions_before) not in remaining
+
+        await store.documents_table.checkout("release-1")
+        rows = await store.documents_table.count_rows()
+        await store.documents_table.checkout_latest()
+        assert rows == 2
+
+
 @pytest.mark.asyncio
 async def test_tag_writes_raise_when_read_only(temp_db_path):
     async with Store(temp_db_path, create=True) as store: