Close connections if __aenter__ fails mid-initialization

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -96,7 +96,14 @@ class HaikuRAG:
             read_only=self._read_only,
             before=self._before,
         )
-        await self.store._initialize()
+        # If _initialize fails mid-way (e.g. migration check raises after
+        # connect), close the store so we don't leak the LanceDB connection —
+        # __aexit__ won't run because the `async with` never entered.
+        try:
+            await self.store._initialize()
+        except BaseException:
+            self.store.close()
+            raise
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
         self.document_item_repository = DocumentItemRepository(self.store)

haiku_rag_slim/haiku/rag/client/rebuild.py

@@ -29,7 +29,10 @@ async def rebuild_database(
     if mode is None:
         mode = RebuildMode.FULL
 
-    # Wait for any background vacuum before destructive table operations
+    # Wait for any background vacuum before destructive table operations.
+    # Rebuild drops and recreates tables (+ creates indices); a concurrent
+    # optimize on the same table fails with "CreateIndex transaction was
+    # preempted" from lance.
     await client._await_vacuum_tasks()
 
     # Update settings to current config

haiku_rag_slim/haiku/rag/store/engine.py

@@ -274,7 +274,14 @@ class Store:
             await self._validate_configuration()
 
     async def __aenter__(self):
-        await self._initialize()
+        # If _initialize connects to LanceDB but then fails (e.g. migration
+        # check, config validation), close the connection so it doesn't
+        # leak — __aexit__ won't run because the `async with` never entered.
+        try:
+            await self._initialize()
+        except BaseException:
+            self.close()
+            raise
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002

tests/test_lancedb_connection.py

@@ -238,3 +238,62 @@ class TestStoreSkipsPathValidationForRemote:
                     skip_migration_check=True,
                 ) as store:
                     assert store is not None
+
+
+class TestInitFailureCleanup:
+    @pytest.mark.asyncio
+    async def test_store_aenter_closes_connection_on_init_failure(
+        self, temp_db_path, monkeypatch
+    ):
+        """If _initialize raises after connect, __aenter__ must close the
+        AsyncConnection so it doesn't leak (no __aexit__ runs in that case)."""
+        mock_conn = AsyncMock()
+        mock_conn.close = lambda: mock_conn.close_calls.append(True)  # type: ignore[attr-defined]
+        mock_conn.close_calls = []  # type: ignore[attr-defined]
+
+        async def fake_connect(*args, **kwargs):
+            return mock_conn
+
+        async def failing_init_tables(self):
+            raise RuntimeError("simulated table init failure")
+
+        monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
+        monkeypatch.setattr(Store, "_init_tables", failing_init_tables)
+
+        with pytest.raises(RuntimeError, match="simulated table init failure"):
+            async with Store(temp_db_path, create=True) as store:
+                assert store is not None
+
+        assert mock_conn.close_calls == [True], (
+            "AsyncConnection.close() was not called on init failure"
+        )
+
+    @pytest.mark.asyncio
+    async def test_client_aenter_closes_store_on_init_failure(
+        self, temp_db_path, monkeypatch
+    ):
+        """HaikuRAG.__aenter__ must close the store if _initialize fails."""
+        from haiku.rag.client import HaikuRAG
+
+        close_calls: list[bool] = []
+
+        original_close = Store.close
+
+        def tracking_close(self):
+            close_calls.append(True)
+            original_close(self)
+
+        async def failing_init(self):
+            # Set db so close() has something to close
+            self.db = AsyncMock()
+            self.db.close = lambda: None
+            raise RuntimeError("simulated initialize failure")
+
+        monkeypatch.setattr(Store, "_initialize", failing_init)
+        monkeypatch.setattr(Store, "close", tracking_close)
+
+        with pytest.raises(RuntimeError, match="simulated initialize failure"):
+            async with HaikuRAG(temp_db_path, create=True):
+                pass
+
+        assert close_calls, "Store.close() was not called when _initialize raised"