Make sure vacuum task has been complete before exiting the client

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
 

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 [

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()