fix(indexer): roll back the PageIndex blob when indexing fails after col.add

Since the add mutation stopped eagerly snapshotting .openkb/files (it registers the new blob via track_new only on success), a blob written by col.add() leaks if index_long_document raises afterward: pageindex.db is rolled back by the snapshot but the blob file is not, and nothing reclaims it. Delete the doc on the failure path so the indexer owns cleanup of a half-applied add, restoring the .openkb/files rollback-surface guarantee.

Author: gwokhou

openkb/indexer.py

@@ -188,49 +188,68 @@ def index_long_document(
             if attempt == max_retries:
                 raise RuntimeError(f"Failed to index {pdf_path.name} after {max_retries} attempts: {exc}") from exc
 
-    # Fetch complete document (metadata + structure + text)
-    doc = col.get_document(doc_id, include_text=True)
-    indexed_doc_name: str = doc.get("doc_name", pdf_path.stem)
-    description: str = doc.get("doc_description", "")
-    structure: list = doc.get("structure", [])
-
-    # Debug: print doc keys and page_count to diagnose get_page_content range
-    logger.info("Doc keys: %s", list(doc.keys()))
-    logger.info("page_count from doc: %s", doc.get("page_count", "NOT PRESENT"))
-
-    tree = {
-        "doc_name": indexed_doc_name,
-        "doc_description": description,
-        "structure": structure,
-    }
-
-    # Write wiki/sources/ — per-page content
-    sources_dir = kb_dir / "wiki" / "sources"
-    sources_dir.mkdir(parents=True, exist_ok=True)
-    images_dir = sources_dir / "images" / source_name
-
-    all_pages: list[dict[str, Any]] = []
-    if pageindex_api_key:
-        # Cloud mode: fetch OCR'd markdown from PageIndex. get_page_content
-        # requires a page range, so pass "1-N".
-        page_count = _get_pdf_page_count(pdf_path)
-        try:
-            all_pages = _normalize_page_content(col.get_page_content(doc_id, f"1-{page_count}"))
-        except Exception as exc:
-            logger.warning("Cloud get_page_content failed for %s: %s", pdf_path.name, exc)
-
-    if not all_pages:
+    # The PageIndex blob for doc_id is now durably on disk. The add mutation no
+    # longer eagerly snapshots .openkb/files — it registers the new blob via
+    # snapshot.track_new() only on a successful return — so if any step below
+    # fails, delete the document we just added. Otherwise the blob leaks as an
+    # orphan that pageindex.db (rolled back by the snapshot) no longer refs and
+    # no reaper reclaims.
+    try:
+        # Fetch complete document (metadata + structure + text)
+        doc = col.get_document(doc_id, include_text=True)
+        indexed_doc_name: str = doc.get("doc_name", pdf_path.stem)
+        description: str = doc.get("doc_description", "")
+        structure: list = doc.get("structure", [])
+
+        # Debug: print doc keys and page_count to diagnose get_page_content range
+        logger.info("Doc keys: %s", list(doc.keys()))
+        logger.info("page_count from doc: %s", doc.get("page_count", "NOT PRESENT"))
+
+        tree = {
+            "doc_name": indexed_doc_name,
+            "doc_description": description,
+            "structure": structure,
+        }
+
+        # Write wiki/sources/ — per-page content
+        sources_dir = kb_dir / "wiki" / "sources"
+        sources_dir.mkdir(parents=True, exist_ok=True)
+        images_dir = sources_dir / "images" / source_name
+
+        all_pages: list[dict[str, Any]] = []
         if pageindex_api_key:
-            logger.warning("Cloud returned no pages for %s; falling back to local pymupdf", pdf_path.name)
-        all_pages = _normalize_page_content(_convert_pdf_to_pages(pdf_path, source_name, images_dir))
-
-    if not all_pages:
-        raise RuntimeError(f"No page content extracted for {pdf_path.name}")
-
-    _write_long_doc_artifacts(
-        tree, all_pages, source_name, doc_id, kb_dir, description=description
-    )
-    return IndexResult(doc_id=doc_id, description=description, tree=tree)
+            # Cloud mode: fetch OCR'd markdown from PageIndex. get_page_content
+            # requires a page range, so pass "1-N".
+            page_count = _get_pdf_page_count(pdf_path)
+            try:
+                all_pages = _normalize_page_content(col.get_page_content(doc_id, f"1-{page_count}"))
+            except Exception as exc:
+                logger.warning("Cloud get_page_content failed for %s: %s", pdf_path.name, exc)
+
+        if not all_pages:
+            if pageindex_api_key:
+                logger.warning("Cloud returned no pages for %s; falling back to local pymupdf", pdf_path.name)
+            all_pages = _normalize_page_content(_convert_pdf_to_pages(pdf_path, source_name, images_dir))
+
+        if not all_pages:
+            raise RuntimeError(f"No page content extracted for {pdf_path.name}")
+
+        _write_long_doc_artifacts(
+            tree, all_pages, source_name, doc_id, kb_dir, description=description
+        )
+        return IndexResult(doc_id=doc_id, description=description, tree=tree)
+    except BaseException:
+        # Best-effort: remove the blob this add created. A failure here (e.g. a
+        # second interrupt) only means the blob may stay orphaned — the original
+        # error still propagates so the caller (mutation coordinator) rolls back
+        # everything else it snapshotted.
+        try:
+            col.delete_document(doc_id)
+        except Exception:
+            logger.warning(
+                "PageIndex cleanup of %s failed after error; blob may be orphaned", doc_id
+            )
+        raise
 
 
 # PageIndex's get_page_content rejects a single page range covering more than

tests/test_indexer.py

@@ -89,6 +89,29 @@ def test_returns_index_result(self, kb_dir, sample_tree, tmp_path):
         assert result.description == sample_tree["doc_description"]
         assert result.tree is not None
 
+    def test_deletes_pageindex_doc_when_a_post_add_step_fails(self, kb_dir, sample_tree, tmp_path):
+        """The PageIndex blob is durably written by col.add(), but .openkb/files is
+        no longer in the add mutation's eager snapshot — track_new only registers
+        the blob on a successful return. So if any step after col.add() raises
+        (here: get_document), index_long_document must delete the doc it just
+        added; otherwise the blob leaks as an orphan that pageindex.db — rolled
+        back by the snapshot — no longer references, and no reaper reclaims."""
+        doc_id = "abc-123"
+        col = self._make_fake_collection(doc_id, sample_tree)
+        col.get_document.side_effect = RuntimeError("get_document blew up")
+
+        fake_client = MagicMock()
+        fake_client.collection.return_value = col
+
+        pdf_path = tmp_path / "sample.pdf"
+        pdf_path.write_bytes(b"%PDF-1.4 fake")
+
+        with patch("openkb.indexer.PageIndexClient", return_value=fake_client):
+            with pytest.raises(RuntimeError, match="get_document blew up"):
+                index_long_document(pdf_path, kb_dir)
+
+        col.delete_document.assert_called_once_with(doc_id)
+
     def test_source_page_written_as_json(self, kb_dir, sample_tree, tmp_path):
         """Long doc source should be written as JSON, not markdown."""
         import json as json_mod