add deadlock mitigations for vacuum

haiku_rag_slim/haiku/rag/config/models.py

@@ -59,6 +59,11 @@ class StorageConfig(BaseModel):
     data_dir: Path = Field(default_factory=get_default_data_dir)
     auto_vacuum: bool = True
     vacuum_retention_seconds: int = 86400
+    # Upper bound (seconds) on a single vacuum/optimize pass. ``None`` leaves it
+    # unbounded (the historical behavior). A finite value converts a stuck
+    # compaction into a logged, skipped pass instead of an unbounded hang that
+    # would also wedge client teardown (which drains background vacuums).
+    vacuum_timeout_seconds: float | None = None
 
 
 class LanceDBConfig(BaseModel):

haiku_rag_slim/haiku/rag/store/engine.py

@@ -571,14 +571,29 @@ class Store:
                 # Evaluate config at runtime to allow dynamic changes
                 if retention_seconds is None:
                     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._tables().values():
-                    await table.optimize(
-                        cleanup_older_than=await self._tag_safe_retention(
-                            table, retention
+                # Bound the pass so a stuck compaction degrades to a logged,
+                # skipped pass rather than an unbounded await -- which would also
+                # wedge client teardown (__aexit__ drains background vacuums).
+                # timeout=None leaves it unbounded (historical behavior).
+                timeout = self._config.storage.vacuum_timeout_seconds
+                async with asyncio.timeout(timeout):
+                    # Perform maintenance per table using optimize() with configurable retention
+                    retention = timedelta(seconds=retention_seconds)
+                    for table in self._tables().values():
+                        await table.optimize(
+                            cleanup_older_than=await self._tag_safe_retention(
+                                table, retention
+                            )
                         )
-                    )
+            except TimeoutError:
+                # The pass is best-effort maintenance; a timeout is not fatal.
+                # NOTE: cancellation only takes effect if the underlying
+                # operation is cancellable -- a fully-unresponsive native future
+                # may still linger, but the caller's coroutine is freed.
+                logger.warning(
+                    "Vacuum timed out after %ss; skipping this pass",
+                    self._config.storage.vacuum_timeout_seconds,
+                )
             except OSError as e:
                 # Resource errors (e.g. disk pressure) skip the pass; lance
                 # errors surface as RuntimeError and must not be swallowed —

tests/test_vacuum_debounce.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 
 import pytest
 
@@ -18,6 +19,58 @@ def _docling_doc(name: str, text: str):
     return doc
 
 
+@pytest.mark.asyncio
+async def test_vacuum_times_out_and_skips(temp_db_path, monkeypatch, caplog):
+    """A stuck optimize is bounded by storage.vacuum_timeout_seconds: the pass
+    is skipped (logged), the call returns instead of hanging, and no exception
+    propagates."""
+
+    class _SlowTable:
+        async def optimize(self, **_kwargs):
+            await asyncio.sleep(30)  # far longer than the timeout
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        store = client.store
+        # setattr (not direct assignment) so the shared global Config is restored.
+        monkeypatch.setattr(store._config.storage, "vacuum_timeout_seconds", 0.05)
+        monkeypatch.setattr(store, "_tables", lambda: {"t": _SlowTable()})
+
+        async def _passthrough_retention(_table, retention):
+            return retention
+
+        monkeypatch.setattr(store, "_tag_safe_retention", _passthrough_retention)
+
+        with caplog.at_level(logging.WARNING, logger="haiku.rag.store.engine"):
+            # wait_for guards the test itself from hanging if the bound fails.
+            await asyncio.wait_for(store.vacuum(), timeout=5)
+
+        assert "Vacuum timed out" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_vacuum_unbounded_when_timeout_none(temp_db_path, monkeypatch):
+    """timeout=None preserves the historical unbounded behavior (optimize runs
+    to completion)."""
+    ran: list[int] = []
+
+    class _Table:
+        async def optimize(self, **_kwargs):
+            ran.append(1)
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        store = client.store
+        monkeypatch.setattr(store._config.storage, "vacuum_timeout_seconds", None)
+        monkeypatch.setattr(store, "_tables", lambda: {"t": _Table()})
+
+        async def _passthrough_retention(_table, retention):
+            return retention
+
+        monkeypatch.setattr(store, "_tag_safe_retention", _passthrough_retention)
+
+        await store.vacuum()
+        assert ran == [1]
+
+
 @pytest.mark.asyncio
 async def test_schedule_vacuum_is_debounced(temp_db_path, monkeypatch):
     """Rapid writes within the throttle window schedule only one background