Improved Document Indexing/ Search

src/holoviz_mcp/holoviz_mcp/data.py

@@ -4,6 +4,8 @@
 import logging
 import os
 import re
+import shutil
+import time
 from pathlib import Path
 from typing import Any
 from typing import Literal
@@ -12,6 +14,7 @@
 import chromadb
 import git
 from chromadb.api.collection_configuration import CreateCollectionConfiguration
+from chromadb.api.shared_system_client import SharedSystemClient
 from fastmcp import Context
 from nbconvert import MarkdownExporter
 from nbformat import read as nbread
@@ -490,10 +493,20 @@ def __init__(self, *, data_dir: Optional[Path] = None, repos_dir: Optional[Path]
         if not config.server.anonymized_telemetry:
             os.environ["ANONYMIZED_TELEMETRY"] = "False"
 
-        # Initialize ChromaDB
-        self.chroma_client = chromadb.PersistentClient(path=str(self._vector_db_path))
-
-        self.collection = self.chroma_client.get_or_create_collection("holoviz_docs", configuration=_CROMA_CONFIGURATION)
+        # Initialize ChromaDB with health check
+        try:
+            self.chroma_client = chromadb.PersistentClient(path=str(self._vector_db_path))
+            self.collection = self.chroma_client.get_or_create_collection("holoviz_docs", configuration=_CROMA_CONFIGURATION)
+            self.collection.count()  # Probe read to verify health
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except BaseException as e:
+            logger.error("ChromaDB initialization failed (%s: %s). Wiping and reinitializing.", type(e).__name__, e)
+            shutil.rmtree(self._vector_db_path, ignore_errors=True)
+            self._vector_db_path.mkdir(parents=True, exist_ok=True)
+            SharedSystemClient.clear_system_cache()
+            self.chroma_client = chromadb.PersistentClient(path=str(self._vector_db_path))
+            self.collection = self.chroma_client.get_or_create_collection("holoviz_docs", configuration=_CROMA_CONFIGURATION)
 
         # Lazy-initialized async lock for database operations to prevent corruption from concurrent access
         self._db_lock: Optional[asyncio.Lock] = None
@@ -527,12 +540,37 @@ def db_lock(self) -> asyncio.Lock:
             self._db_lock = asyncio.Lock()
         return self._db_lock
 
+    @property
+    def _backup_path(self) -> Path:
+        """Path to the backup copy of the vector database directory."""
+        return Path(str(self._vector_db_path) + ".bak")
+
+    async def _restore_from_backup(self, ctx: Context | None = None) -> None:
+        """Restore vector database from backup after a write failure."""
+        backup_path = self._backup_path
+        if not backup_path.exists():
+            await log_warning("No backup found to restore from.", ctx)
+            return
+        try:
+            shutil.rmtree(self._vector_db_path, ignore_errors=True)
+            shutil.copytree(backup_path, self._vector_db_path)
+            SharedSystemClient.clear_system_cache()
+            self.chroma_client = chromadb.PersistentClient(path=str(self._vector_db_path))
+            self.collection = self.chroma_client.get_or_create_collection("holoviz_docs", configuration=_CROMA_CONFIGURATION)
+            await log_info("Restored vector database from backup.", ctx)
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except BaseException as e:
+            logger.error("Failed to restore from backup (%s: %s). Database may be degraded.", type(e).__name__, e)
+
     def is_indexed(self) -> bool:
         """Check if documentation index exists and is valid."""
         try:
             count = self.collection.count()
             return count > 0
-        except Exception:
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except BaseException:
             return False
 
     async def ensure_indexed(self, ctx: Context | None = None):
@@ -887,6 +925,16 @@ async def index_documentation(self, ctx: Context | None = None):
         # Validate for duplicate IDs and log details
         await self._validate_unique_ids(all_docs)
 
+        # Create pre-write backup of vector database
+        backup_path = self._backup_path
+        try:
+            if self._vector_db_path.exists():
+                t0 = time.monotonic()
+                shutil.copytree(self._vector_db_path, backup_path, dirs_exist_ok=True)
+                await log_info(f"Pre-write backup created at {backup_path} ({time.monotonic() - t0:.1f}s)", ctx)
+        except Exception as e:
+            await log_warning(f"Failed to create pre-write backup: {e}. Continuing without backup.", ctx)
+
         # Clear existing collection
         await log_info("Clearing existing index...", ctx)
 
@@ -911,23 +959,30 @@ async def index_documentation(self, ctx: Context | None = None):
         # Add documents to ChromaDB
         await log_info(f"Adding {len(all_docs)} documents to index...", ctx)
 
-        self.collection.add(
-            documents=[doc["content"] for doc in all_docs],
-            metadatas=[
-                {
-                    "title": doc["title"],
-                    "url": doc["url"],
-                    "project": doc["project"],
-                    "source_path": doc["source_path"],
-                    "source_path_stem": doc["source_path_stem"],
-                    "source_url": doc["source_url"],
-                    "description": doc["description"],
-                    "is_reference": doc["is_reference"],
-                }
-                for doc in all_docs
-            ],
-            ids=[doc["id"] for doc in all_docs],
-        )
+        try:
+            self.collection.add(
+                documents=[doc["content"] for doc in all_docs],
+                metadatas=[
+                    {
+                        "title": doc["title"],
+                        "url": doc["url"],
+                        "project": doc["project"],
+                        "source_path": doc["source_path"],
+                        "source_path_stem": doc["source_path_stem"],
+                        "source_url": doc["source_url"],
+                        "description": doc["description"],
+                        "is_reference": doc["is_reference"],
+                    }
+                    for doc in all_docs
+                ],
+                ids=[doc["id"] for doc in all_docs],
+            )
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except BaseException as e:
+            logger.error("ChromaDB write failed (%s: %s). Attempting restore from backup.", type(e).__name__, e)
+            await self._restore_from_backup(ctx)
+            raise
 
         await log_info(f"✅ Successfully indexed {len(all_docs)} documents", ctx)
         await log_info(f"📊 Vector database stored at: {self._vector_db_path}", ctx)

tests/docs_mcp/test_data.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+import chromadb
 import pytest
 from pydantic import AnyHttpUrl
 
@@ -472,3 +473,106 @@ def test_truncate_content_with_special_characters():
     # Should handle special characters gracefully
     assert result is not None
     assert "foo-bar" in result or "baz_qux" in result or "Start text" in result
+
+
+# --- ChromaDB health check and backup tests ---
+
+
+def test_init_health_check_recovers_from_corrupt_db(tmp_path):
+    """ChromaDB init recovers from a corrupt database by wiping and reinitializing."""
+    from chromadb.api.shared_system_client import SharedSystemClient
+
+    vector_dir = tmp_path / "chroma"
+    vector_dir.mkdir()
+
+    # Write a corrupt chroma.sqlite3 to trigger an init failure
+    (vector_dir / "chroma.sqlite3").write_text("THIS IS NOT A VALID SQLITE FILE")
+
+    # Clear cached ChromaDB state from other tests to avoid stale references
+    SharedSystemClient.clear_system_cache()
+
+    indexer = DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=vector_dir)
+
+    # Should have recovered: collection exists and is empty
+    assert indexer.collection.count() == 0
+
+
+def test_init_health_check_reraises_keyboard_interrupt(tmp_path):
+    """KeyboardInterrupt during ChromaDB init is not swallowed."""
+    from unittest.mock import patch
+
+    vector_dir = tmp_path / "chroma"
+    vector_dir.mkdir()
+
+    with patch("chromadb.PersistentClient", side_effect=KeyboardInterrupt):
+        with pytest.raises(KeyboardInterrupt):
+            DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=vector_dir)
+
+
+def test_is_indexed_catches_base_exception(tmp_path):
+    """is_indexed() returns False when collection.count() raises a BaseException subclass."""
+    from unittest.mock import PropertyMock
+    from unittest.mock import patch
+
+    indexer = DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=tmp_path / "chroma")
+
+    # Simulate a PanicException (BaseException subclass) from ChromaDB's Rust backend
+    class FakePanicException(BaseException):
+        pass
+
+    with patch.object(type(indexer.collection), "count", new_callable=PropertyMock) as mock_count:
+        mock_count.return_value = FakePanicException("rust panic")
+        # Replace count as a regular method that raises
+        indexer.collection.count = lambda: (_ for _ in ()).throw(FakePanicException("rust panic"))  # type: ignore[assignment]
+        assert indexer.is_indexed() is False
+
+
+def test_backup_path_property(tmp_path):
+    """_backup_path returns the correct sibling path with .bak suffix."""
+    indexer = DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=tmp_path / "chroma")
+    assert indexer._backup_path == tmp_path / "chroma.bak"
+
+
+@pytest.mark.asyncio
+async def test_restore_from_backup_on_write_failure(tmp_path):
+    """On write failure, _restore_from_backup recovers original data."""
+    import shutil
+
+    from chromadb.api.shared_system_client import SharedSystemClient
+
+    vector_dir = tmp_path / "chroma"
+    SharedSystemClient.clear_system_cache()
+    indexer = DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=vector_dir)
+
+    # Seed the collection with one document
+    indexer.collection.add(documents=["hello world"], metadatas=[{"project": "test"}], ids=["doc1"])
+    assert indexer.collection.count() == 1
+
+    # Force-flush by recreating the client so all WAL data is checkpointed
+    del indexer.chroma_client
+    SharedSystemClient.clear_system_cache()
+    client = chromadb.PersistentClient(path=str(vector_dir))
+    assert client.get_or_create_collection("holoviz_docs").count() == 1
+    del client
+    SharedSystemClient.clear_system_cache()
+
+    # Create a backup (simulating what index_documentation does)
+    backup_path = indexer._backup_path
+    shutil.copytree(vector_dir, backup_path, dirs_exist_ok=True)
+
+    # Recreate indexer to get fresh client after backup
+    indexer = DocumentationIndexer(data_dir=tmp_path, repos_dir=tmp_path / "repos", vector_dir=vector_dir)
+
+    # Now corrupt the live database by clearing and adding a bad doc
+    indexer.collection.delete(ids=["doc1"])
+    indexer.collection.add(documents=["corrupted"], metadatas=[{"project": "bad"}], ids=["bad1"])
+    assert indexer.collection.count() == 1
+
+    # Restore from backup
+    await indexer._restore_from_backup()
+
+    # After restore, should have the original document back
+    assert indexer.collection.count() == 1
+    result = indexer.collection.get(ids=["doc1"])
+    assert result["ids"] == ["doc1"]
+    assert result["documents"] == ["hello world"]