fix(cloud-import): stop windowed page fetch on a short window, not the tree's max index

_fetch_cloud_pages bounded its loop by the tree's max page index
(_max_page_index). When the cloud tree under-reports the page count — a real
case (e.g. a paper whose tree stops a couple pages short of the references) — the
loop exited before fetching a later window and silently DROPPED pages. At the
exact 1000 boundary this was an off-by-one (a 1001-page doc whose tree maxes at
1000 lost page 1001); more generally any >1000-page doc with an under-reporting
tree was truncated.

Drop the max-index bound: request fixed 1000-page windows and stop as soon as a
window comes back SHORT (PageIndex page numbers are sequential, so a short window
means we've passed the last page). The common ≤1000-page doc stays a single
request, a larger doc fetches every page, and an under-reported tree no longer
truncates. Remove the now-unused _max_page_index. Tests updated + a
full-window-triggers-next-fetch regression.

Author: KylinMountain

openkb/indexer.py

@@ -213,55 +213,35 @@ def index_long_document(
 # this many pages (``parse_pages`` raises "Page range too large (max 1000)"),
 # so cloud page fetches are windowed in chunks of this size.
 _CLOUD_PAGE_WINDOW = 1000
+# Safety bound on the windowed fetch (in pages) in case a backend never returns
+# a short window — caps the loop at _CLOUD_PAGE_MAX / _CLOUD_PAGE_WINDOW calls.
+_CLOUD_PAGE_MAX = 1_000_000
 
 
-def _max_page_index(structure: list, default: int = 0) -> int:
-    """Largest start/end page index across the (possibly nested) tree.
-
-    Cloud tree nodes carry ``start_index``/``end_index`` page numbers; this
-    bounds the windowed cloud page fetch. Returns ``default`` (``0`` = unknown)
-    when no integer indices are present, leaving the fetch to stop on the first
-    empty window.
-    """
-    best = 0
-
-    def _walk(nodes: list) -> None:
-        nonlocal best
-        for node in nodes or []:
-            if not isinstance(node, dict):
-                continue
-            for key in ("end_index", "start_index"):
-                val = node.get(key)
-                if isinstance(val, int):
-                    best = max(best, val)
-            _walk(node.get("nodes"))
-
-    _walk(structure)
-    return best or default
-
-
-def _fetch_cloud_pages(col, doc_id: str, max_page: int) -> list[dict[str, Any]]:
+def _fetch_cloud_pages(col, doc_id: str) -> list[dict[str, Any]]:
     """Fetch all OCR pages of a cloud doc, windowing around the 1000-page cap.
 
     ``get_page_content`` returns the whole document and uses its ``pages`` arg
     only as a client-side filter that ``parse_pages`` caps at 1000 pages — so a
-    single ``"1-<N>"`` request fails for any doc over 1000 pages or whose tree
-    exposes no integer indices. Request fixed ``1000``-page windows instead and
-    concatenate; each window over-covers its range, so an off-by-one or 0-based
-    ``max_page`` never truncates the last page. ``max_page`` (0 = unknown) bounds
-    the loop; otherwise it stops at the first empty window. A wide safety bound
-    prevents an unbounded loop if the backend never returns an empty window.
+    single ``"1-<N>"`` request fails for any doc over 1000 pages. Request fixed
+    ``1000``-page windows and stop as soon as a window comes back SHORT (fewer
+    than a full window): PageIndex page numbers are sequential, so a short window
+    means we've passed the last page. This is what makes the common (≤1000-page)
+    doc a single request, while still fetching every page of a larger one — and,
+    unlike bounding the loop by the tree's max page index, it never truncates a
+    doc whose tree under-reports its page count (a real case: a paper whose tree
+    stops a couple pages short of the references). A wide safety bound guards
+    against a backend that never narrows the window.
     """
     pages: list[dict[str, Any]] = []
-    upper = max_page if max_page and max_page > 0 else 1_000_000
     start = 1
-    while start <= upper:
+    while start <= _CLOUD_PAGE_MAX:
         window = _normalize_page_content(
             col.get_page_content(doc_id, f"{start}-{start + _CLOUD_PAGE_WINDOW - 1}")
         )
-        if not window:
-            break
         pages.extend(window)
+        if len(window) < _CLOUD_PAGE_WINDOW:
+            break
         start += _CLOUD_PAGE_WINDOW
     return pages
 
@@ -303,7 +283,7 @@ def import_cloud_document(doc_id: str, kb_dir: Path, path_key: str) -> CloudImpo
         "structure": structure,
     }
 
-    all_pages = _fetch_cloud_pages(col, doc_id, _max_page_index(structure))
+    all_pages = _fetch_cloud_pages(col, doc_id)
     if not all_pages:
         raise RuntimeError(
             f"No page content returned from PageIndex Cloud for doc_id={doc_id}"

tests/test_indexer.py

@@ -357,17 +357,40 @@ def fake_get(doc_id, rng):
     col = MagicMock()
     col.get_page_content.side_effect = fake_get
 
-    pages = _fetch_cloud_pages(col, "doc", 0)  # 0 = unknown bound → loop until empty
+    pages = _fetch_cloud_pages(col, "doc")
     assert len(pages) == 1500
     assert pages[0]["page"] == 1 and pages[-1]["page"] == 1500
     ranges = [c.args[1] for c in col.get_page_content.call_args_list]
-    assert ranges == ["1-1000", "1001-2000", "2001-3000"]
+    # Full first window → fetch the next; the short 2nd window (500<1000) stops it.
+    assert ranges == ["1-1000", "1001-2000"]
     # Every requested window spans exactly 1000 pages → parse_pages never raises.
     for r in ranges:
         a, b = (int(x) for x in r.split("-"))
         assert b - a + 1 == 1000
 
 
+def test_fetch_cloud_pages_full_window_triggers_next_fetch():
+    """A doc whose pages exactly fill the first window must still fetch the next
+    one. Regression: bounding the loop by the tree's max page index dropped the
+    straggler page(s) of a doc whose tree under-reported its page count.
+    """
+    from openkb.indexer import _fetch_cloud_pages
+
+    def fake_get(doc_id, rng):
+        start = int(rng.split("-")[0])
+        if start == 1:
+            return [{"page": p, "content": "x"} for p in range(1, 1001)]  # full window
+        if start == 1001:
+            return [{"page": 1001, "content": "x"}]  # one straggler past the window
+        return []
+
+    col = MagicMock()
+    col.get_page_content.side_effect = fake_get
+
+    pages = _fetch_cloud_pages(col, "doc")
+    assert [p["page"] for p in pages] == list(range(1, 1002))  # page 1001 NOT dropped
+
+
 def test_import_cloud_document_no_indices_avoids_oversized_range(kb_dir, monkeypatch):
     """A cloud tree with no integer page indices must NOT request a 100000-page
     range (parse_pages rejects >1000); it windows from page 1 instead.