From 5e4928ccadde70858d8a07df3b0c74aa2a7a90ea Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 7 Oct 2025 12:27:55 +0300
Subject: [PATCH] Make sure vacuum task has been complete before exiting the
 client

---
 src/haiku/rag/client.py       |  3 +++
 src/haiku/rag/store/engine.py | 10 ++++++----
 tests/test_versioning.py      | 31 +++++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py
index 97e6f931..e4107b86 100644
--- a/src/haiku/rag/client.py
+++ b/src/haiku/rag/client.py
@@ -46,6 +46,9 @@ class HaikuRAG:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
         """Async context manager exit."""
+        # Wait for any pending vacuum to complete before closing
+        async with self.store._vacuum_lock:
+            pass
         self.close()
         return False
 
diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py
index 9c6b2e88..558f8c0d 100644
--- a/src/haiku/rag/store/engine.py
+++ b/src/haiku/rag/store/engine.py
@@ -80,17 +80,16 @@ class Store:
         if not skip_validation:
             self._validate_configuration()
 
-    async def vacuum(
-        self, retention_seconds: int = Config.VACUUM_RETENTION_SECONDS
-    ) -> None:
+    async def vacuum(self, retention_seconds: int | None = None) -> None:
         """Optimize and clean up old versions across all tables to reduce disk usage.
 
         Args:
             retention_seconds: Retention threshold in seconds. Only versions older
-                              than this will be removed. Defaults to Config.VACUUM_RETENTION_SECONDS.
+                              than this will be removed. If None, uses Config.VACUUM_RETENTION_SECONDS.
 
         Note:
             If vacuum is already running, this method returns immediately without blocking.
+            Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
         """
         if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
             return
@@ -101,6 +100,9 @@ class Store:
 
         async with self._vacuum_lock:
             try:
+                # Evaluate config at runtime to allow dynamic changes
+                if retention_seconds is None:
+                    retention_seconds = Config.VACUUM_RETENTION_SECONDS
                 # Perform maintenance per table using optimize() with configurable retention
                 retention = timedelta(seconds=retention_seconds)
                 for table in [
diff --git a/tests/test_versioning.py b/tests/test_versioning.py
index 869aea8e..874696cd 100644
--- a/tests/test_versioning.py
+++ b/tests/test_versioning.py
@@ -194,3 +194,34 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
     assert after_zero_chunk_versions < initial_chunk_versions, (
         "Should have fewer versions after vacuum(0)"
     )
+
+
+@pytest.mark.asyncio
+async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
+    """Test that background vacuum completes when context manager exits."""
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.config import Config
+    from haiku.rag.utils import text_to_docling_document
+
+    # Set aggressive vacuum retention for this test
+    monkeypatch.setattr(Config, "VACUUM_RETENTION_SECONDS", 0)
+
+    async with HaikuRAG(db_path=temp_db_path) as client:
+        # Create multiple documents - each creation triggers automatic vacuum with retention=0
+        # This aggressively cleans up old versions between operations
+        for i in range(3):
+            doc = Document(content=f"Test document {i}")
+            dl_doc = text_to_docling_document(f"Test document {i}", name=f"test{i}.md")
+            await client.document_repository._create_with_docling(doc, dl_doc)
+
+    # After context exit, automatic vacuum should have kept versions minimal
+    store = Store(temp_db_path)
+    final_versions = len(list(store.documents_table.list_versions()))
+
+    # With retention_seconds=0, vacuum aggressively cleans up between operations
+    # Should have very few versions remaining (1-2)
+    assert final_versions <= 2, (
+        f"Aggressive vacuum should keep minimal versions, got {final_versions}"
+    )
+    assert final_versions >= 1, "Should have at least one version remaining"
+    store.close()