fix: CMYK image drop, empty-doc crash, page_index shadowing, sqlite hardening, flaky tests

Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272.

- pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without-
  alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png
  then raises "unsupported colorspace", silently dropped by the surrounding
  except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves
  RGBA untouched).

- pipeline.py: detect_strategy([]) (an empty/whitespace-only source file)
  returned "content_based", routing into the PDF-oriented TOC-detection
  pipeline -- wasting a real LLM call before raising IndexingError. Empty
  node lists now route to level_based, whose build_tree_from_levels([])
  returns an empty structure instantly with zero LLM calls.

- page_index.py (shim): pageindex/__init__.py binds the canonical `page_index`
  function as the package attribute, but this file is ALSO a real submodule
  of the same name -- importing it anywhere (import machinery, unconditional)
  overwrites that attribute with the module object, breaking
  `from pageindex import page_index; page_index(x)` for the rest of the
  process. Made the shim module itself callable (delegates to the real
  function via a ModuleType subclass), so whichever object ends up in that
  slot is callable regardless of import order.

- storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape
  on a duplicate name (new CollectionAlreadyExistsError); the collections
  table's CHECK constraint only validated the name's first character (GLOB
  '*' is a wildcard, not a regex quantifier over the preceding class) --
  fixed to validate the whole string, and SQLiteStorage now also validates in
  Python (it's a public StorageEngine usable directly, bypassing
  LocalBackend's own check).

- tests/test_review_fixes_2.py: two tests used a ContentNode with no `level`
  set, so build_index took the content_based path and made real (retried,
  slow, and -- with a valid key -- billable) LLM calls instead of testing the
  text-stripping logic they claimed to. Mocked out _content_based_pipeline.

- retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of
  the canonical parse_pages/get_pdf_page_content that had already drifted
  (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical
  now, so the legacy pageindex.get_page_content path can't silently regress
  again.

- parser/markdown.py: a leading UTF-8 BOM broke first-header detection
  (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only
  backtick fences were recognized as code blocks, so a '#'-prefixed line
  inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading --
  recognize both fence styles.

- run_pageindex.py: --if-thinning wasn't migrated to the bare-flag +
  legacy-yes/no convention the other four --if-add-* flags got; bare usage
  raised an argparse error and it never went through the shared coercion.

- types.py: DocumentDetail's `structure` field was inside the class's
  total=False body, so TypedDict rules made it optional even though every
  backend always populates it. Split into a required base class.

Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped
(one pre-existing, unrelated flaky cloud-streaming test).

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS

tests/test_legacy_shims.py

@@ -3,10 +3,15 @@ now deprecation shims over the canonical pageindex.index.* modules. These
 tests pin the compatibility contract."""
 import asyncio
 import importlib
+import subprocess
+import sys
 import warnings
+from pathlib import Path
 
 import pytest
 
+_REPO_ROOT = Path(__file__).resolve().parent.parent
+
 
 def test_plain_import_pageindex_does_not_warn():
     # `import pageindex` must not route through the deprecation shims.
@@ -116,3 +121,27 @@ def test_md_to_tree_coerces_legacy_yes_no_strings(tmp_path):
 
     assert not _has_summary(result["structure"])
     assert all("node_id" in n for n in result["structure"])
+
+
+def test_page_index_stays_callable_after_the_submodule_is_imported():
+    """pageindex/__init__.py binds the FUNCTION `page_index` as the package
+    attribute, but pageindex/page_index.py is ALSO a real submodule of the
+    same name — importing that submodule anywhere clobbers the package
+    attribute with the module object (Python's import machinery does this
+    unconditionally). Must run in a fresh subprocess: the effect depends on
+    import order, so it can't be reliably observed against an
+    already-imported pageindex in this test process."""
+    script = (
+        "import warnings; warnings.simplefilter('ignore')\n"
+        "import pageindex.page_index\n"  # the clobbering import
+        "from pageindex import page_index\n"
+        "assert callable(page_index), f'page_index is not callable: {type(page_index)}'\n"
+        "from pageindex.page_index import page_index_main\n"  # old multi-symbol import still works
+        "assert callable(page_index_main)\n"
+        "print('OK')\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, cwd=str(_REPO_ROOT),
+    )
+    assert result.returncode == 0, result.stderr
+    assert "OK" in result.stdout

tests/test_local_backend.py

@@ -35,6 +35,27 @@ def test_unsupported_file_type_raises(backend, tmp_path):
         backend.add_document("papers", str(bad_file))
 
 
+def test_add_document_on_empty_markdown_file_does_not_crash(tmp_path):
+    """An empty/whitespace-only .md file used to route into the PDF-oriented
+    TOC-detection pipeline (no node ever has 'level' set), wasting an LLM call
+    and then raising IndexingError. Must complete instantly with zero LLM
+    calls when summary/description are off."""
+    from pageindex.config import IndexConfig
+
+    storage = SQLiteStorage(str(tmp_path / "test.db"))
+    backend = LocalBackend(
+        storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o",
+        index_config=IndexConfig(if_add_node_summary=False, if_add_doc_description=False),
+    )
+    backend.get_or_create_collection("papers")
+    empty_md = tmp_path / "empty.md"
+    empty_md.write_text("   \n\n  \n")
+
+    doc_id = backend.add_document("papers", str(empty_md))  # must not raise
+
+    assert backend.get_document_structure("papers", doc_id) == []
+
+
 def test_register_custom_parser(backend):
     from pageindex.parser.protocol import ParsedDocument, ContentNode
 

tests/test_markdown_parser.py

@@ -71,3 +71,32 @@ def test_headerless_file_yields_single_node(tmp_path):
     assert len(result.nodes) == 1
     assert result.nodes[0].title == "plain"
     assert "No headings at all" in result.nodes[0].content
+
+
+def test_utf8_bom_does_not_break_the_first_header(tmp_path):
+    """A leading BOM (common from Windows editors/exporters) isn't
+    whitespace, so .strip() doesn't remove it — without utf-8-sig decoding,
+    the header regex fails to match the BOM-prefixed first line, and it gets
+    misclassified as unrecognized preamble text instead of a real heading."""
+    md = tmp_path / "bom.md"
+    md.write_bytes(b"\xef\xbb\xbf# First Header\nbody text\n")
+    result = MarkdownParser().parse(str(md))
+    assert len(result.nodes) == 1
+    assert result.nodes[0].title == "First Header"
+    assert result.nodes[0].level == 1
+
+
+def test_tilde_fenced_code_blocks_are_recognized(tmp_path):
+    """CommonMark allows both backtick and tilde code fences. Only
+    recognizing backticks let a '#'-prefixed line inside a ~~~-fenced block
+    (e.g. a shell comment in a sample) be misparsed as a real heading."""
+    md = tmp_path / "tilde.md"
+    md.write_text(
+        "# Real Header\nintro\n"
+        "~~~\n# not a real header, just a comment\n~~~\n"
+        "## Real Sub\nmore\n"
+    )
+    result = MarkdownParser().parse(str(md))
+    titles = [n.title for n in result.nodes]
+    assert titles == ["Real Header", "Real Sub"]
+    assert "not a real header" not in " ".join(n.title for n in result.nodes)

tests/test_page_content.py

@@ -34,3 +34,43 @@ def test_retrieve_md_page_content_returns_only_requested_lines():
 
     out = _get_md_page_content({"structure": _md_structure()}, [5, 100])
     assert [r["page"] for r in out] == [5, 100]
+
+
+def test_retrieve_parse_pages_delegates_to_canonical_and_enforces_dos_cap():
+    """retrieve._parse_pages used to be an independent copy that lacked the
+    canonical parse_pages' p>=1 filter and 1000-page cap — a caller of the
+    legacy pageindex.get_page_content could bypass the DoS guard the SDK path
+    enforces. Now it's a one-line delegate, so they can't drift again."""
+    from pageindex.retrieve import _parse_pages
+    from pageindex.index.utils import parse_pages
+    import pytest
+
+    assert _parse_pages("5-7") == parse_pages("5-7") == [5, 6, 7]
+    with pytest.raises(ValueError, match="too large"):
+        _parse_pages("1-99999999")
+
+
+def test_retrieve_get_pdf_page_content_falls_back_to_canonical(tmp_path, monkeypatch):
+    """When no cached 'pages' are present, the file-read fallback must
+    delegate to the canonical get_pdf_page_content instead of re-implementing
+    PDF text extraction inline (a second, independently-maintained copy)."""
+    from pageindex.retrieve import _get_pdf_page_content
+    import pageindex.retrieve as retrieve_mod
+
+    calls = []
+    monkeypatch.setattr(
+        retrieve_mod, "get_pdf_page_content",
+        lambda path, page_nums: calls.append((path, page_nums)) or [{"page": 1, "content": "x"}],
+    )
+    result = _get_pdf_page_content({"path": "/fake/doc.pdf"}, [1])
+    assert calls == [("/fake/doc.pdf", [1])]
+    assert result == [{"page": 1, "content": "x"}]
+
+
+def test_retrieve_get_pdf_page_content_prefers_cache_over_file():
+    from pageindex.retrieve import _get_pdf_page_content
+
+    doc_info = {"path": "/should/not/be/opened.pdf",
+                "pages": [{"page": 1, "content": "cached one"}, {"page": 2, "content": "cached two"}]}
+    result = _get_pdf_page_content(doc_info, [2])
+    assert result == [{"page": 2, "content": "cached two"}]

tests/test_pdf_parser.py

@@ -1,3 +1,4 @@
+import pymupdf
 import pytest
 from pathlib import Path
 from pageindex.parser.pdf import PdfParser
@@ -29,6 +30,24 @@ def test_parse_nodes_are_flat_without_level():
         assert node.level is None
 
 
+def test_cmyk_pixmap_without_alpha_is_saveable_as_png(tmp_path):
+    """A CMYK image with no alpha has n==4 -- same as RGBA -- so `pix.n > 4`
+    wrongly skips the RGB conversion PNG needs, and pix.save() raises
+    'unsupported colorspace for png', silently dropping the image via the
+    extractor's bare except. The fix (`pix.n - pix.alpha >= 4`) must convert
+    CMYK (4-0=4) while leaving RGBA (4-1=3) untouched."""
+    cmyk = pymupdf.Pixmap(pymupdf.csCMYK, pymupdf.Rect(0, 0, 10, 10))
+    assert cmyk.n == 4 and cmyk.alpha == 0
+    assert cmyk.n - cmyk.alpha >= 4  # the fixed condition: must convert
+    converted = pymupdf.Pixmap(pymupdf.csRGB, cmyk)
+    converted.save(str(tmp_path / "cmyk.png"))  # must not raise
+
+    rgba = pymupdf.Pixmap(pymupdf.Pixmap(pymupdf.csRGB, pymupdf.Rect(0, 0, 10, 10)), 1)
+    assert rgba.n == 4 and rgba.alpha == 1
+    assert not (rgba.n - rgba.alpha >= 4)  # unchanged: RGBA needs no conversion
+    rgba.save(str(tmp_path / "rgba.png"))  # already saveable as-is
+
+
 def test_image_paths_are_absolute(tmp_path):
     """Image references must be absolute so they resolve regardless of cwd
     (cwd-relative paths broke after the query ran from another directory)."""

tests/test_pipeline.py

@@ -25,6 +25,22 @@ def test_detect_strategy_without_level():
     assert detect_strategy(nodes) == "content_based"
 
 
+def test_detect_strategy_empty_nodes_is_level_based():
+    """An empty node list (e.g. an empty/whitespace-only source file) must
+    route to level_based, whose build_tree_from_levels([]) returns an empty
+    structure with zero LLM calls — not content_based, whose TOC-detection
+    pipeline needs real page content and wastes an LLM call before failing."""
+    assert detect_strategy([]) == "level_based"
+
+
+def test_build_index_on_empty_document_makes_no_llm_calls():
+    from pageindex.config import IndexConfig
+    parsed = ParsedDocument(doc_name="empty", nodes=[])
+    opt = IndexConfig(if_add_node_summary=False, if_add_doc_description=False)
+    result = build_index(parsed, opt=opt)
+    assert result == {"doc_name": "empty", "structure": []}
+
+
 def test_build_tree_from_levels():
     nodes = [
         ContentNode(content="ch1 text", tokens=10, title="Chapter 1", index=1, level=1),

tests/test_review_fixes_2.py

@@ -2,6 +2,7 @@
 2d46d68..8f536cb): the Markdown text-stripping fix's fallout, plus the other
 directly-fixable findings from that pass."""
 import asyncio
+from unittest.mock import AsyncMock
 
 import pytest
 
@@ -122,6 +123,20 @@ def test_cloud_delete_collection_still_clears_real_folder_id(monkeypatch):
 
 
 # ── #7: remove_structure_text is skipped when text was never added ───────────
+def _mock_content_based_pipeline(monkeypatch, structure):
+    """content_based's real path (_content_based_pipeline) drives real LLM
+    calls (TOC detection etc.) regardless of if_add_node_summary — a prior
+    version of these two tests didn't mock this out, fell through to it, and
+    made real network calls (with a dummy key: 10 retries before failing;
+    with a real key: real billable requests) on every run."""
+    from pageindex.index import pipeline
+
+    async def fake(page_list, opt):
+        return structure
+
+    monkeypatch.setattr(pipeline, "_content_based_pipeline", fake)
+
+
 def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch):
     from pageindex.index import pipeline
     from pageindex.parser.protocol import ContentNode, ParsedDocument
@@ -131,6 +146,7 @@ def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch):
     # ...` inside the function body), so patch it on the utils module itself.
     import pageindex.index.utils as utils_mod
     monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
+    _mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}])
 
     nodes = [ContentNode(content="page one text", tokens=5, index=1)]
     parsed = ParsedDocument(doc_name="d", nodes=nodes)
@@ -147,6 +163,14 @@ def test_build_index_still_strips_text_when_summary_added_it(monkeypatch):
     calls = []
     import pageindex.index.utils as utils_mod
     monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
+    _mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}])
+    # Summary generation itself would otherwise make a real LLM call.
+    monkeypatch.setattr(
+        utils_mod, "generate_summaries_for_structure",
+        AsyncMock(side_effect=lambda structure, model=None: [
+            n.__setitem__("summary", "fake") for n in structure
+        ]),
+    )
 
     nodes = [ContentNode(content="page one text", tokens=5, index=1)]
     parsed = ParsedDocument(doc_name="d", nodes=nodes)

tests/test_sqlite_storage.py

@@ -19,6 +19,60 @@ def test_delete_collection(storage):
     storage.delete_collection("papers")
     assert "papers" not in storage.list_collections()
 
+
+def test_create_duplicate_collection_raises_pageindex_error(storage):
+    """A raw sqlite3.IntegrityError leaking out breaks `except PageIndexError`
+    catch-alls; must be translated to a proper SDK exception."""
+    from pageindex.errors import CollectionAlreadyExistsError, PageIndexError
+    storage.create_collection("papers")
+    with pytest.raises(CollectionAlreadyExistsError):
+        storage.create_collection("papers")
+    # also catchable via the SDK's generic base class
+    storage.create_collection("other")
+    with pytest.raises(PageIndexError):
+        storage.create_collection("other")
+
+
+@pytest.mark.parametrize("bad_name", [
+    "a/../../etc/passwd", "/etc/passwd", "a$(whoami)", ".hidden",
+    "a b", "válid", "", "a" * 129,
+])
+def test_create_collection_rejects_invalid_names_at_the_python_layer(storage, bad_name):
+    """SQLiteStorage must validate collection names itself — it's a public
+    StorageEngine that can be used directly, bypassing LocalBackend's own
+    regex check entirely."""
+    from pageindex.errors import PageIndexError
+    with pytest.raises(PageIndexError):
+        storage.create_collection(bad_name)
+
+
+def test_sql_check_constraint_also_rejects_invalid_names_directly(storage):
+    """Defense-in-depth: even bypassing SQLiteStorage's own Python validation
+    and inserting via raw SQL, the schema's CHECK constraint must reject a
+    name that isn't ENTIRELY [a-zA-Z0-9_-] — not just its first character
+    (GLOB '*' is a wildcard, not a regex quantifier over the preceding class,
+    so 'name GLOB [a-zA-Z0-9_-]*' alone only constrains the first character)."""
+    import sqlite3
+    conn = storage._get_conn()
+    with pytest.raises(sqlite3.IntegrityError):
+        conn.execute("INSERT INTO collections (name) VALUES (?)", ("a/../../etc/passwd",))
+
+
+def test_malicious_collection_name_rejected_through_local_backend_too(tmp_path):
+    """End-to-end via the normal LocalBackend entry point: a path-traversal-
+    shaped collection name must never reach add_document's
+    files_dir / collection path construction. Three independent layers now
+    reject it (LocalBackend's own regex, SQLiteStorage's regex, and the SQL
+    CHECK constraint) — this pins the outermost one."""
+    from pageindex.backend.local import LocalBackend
+    from pageindex.errors import PageIndexError
+
+    storage = SQLiteStorage(str(tmp_path / "t.db"))
+    backend = LocalBackend(storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o")
+    with pytest.raises(PageIndexError):
+        backend.create_collection("a/../../escape_me")
+    assert not (tmp_path / "escape_me").exists()
+
 def test_save_and_get_document(storage):
     storage.create_collection("papers")
     doc = {

tests/test_types.py

@@ -0,0 +1,20 @@
+from pageindex.types import DocumentDetail, DocumentInfo, PageContent
+
+
+def test_document_detail_structure_field_is_required():
+    """structure is always populated by both LocalBackend.get_document and
+    CloudBackend.get_document — must be a required key, not optional, or
+    type checkers/tooling built on this TypedDict wrongly treat a
+    DocumentDetail missing 'structure' as valid."""
+    assert "structure" in DocumentDetail.__required_keys__
+    assert "structure" not in DocumentDetail.__optional_keys__
+
+
+def test_document_detail_backend_specific_fields_stay_optional():
+    assert "file_path" in DocumentDetail.__optional_keys__
+    assert "status" in DocumentDetail.__optional_keys__
+
+
+def test_document_detail_inherits_document_info_as_required():
+    for key in ("doc_id", "doc_name", "doc_description", "doc_type"):
+        assert key in DocumentDetail.__required_keys__