From eeb4e8f85d894d6059181546a55a3e76606be380 Mon Sep 17 00:00:00 2001
From: Research Assistant <research@example.com>
Date: Mon, 1 Jun 2026 11:29:06 +0800
Subject: [PATCH] feat: add paperforge ocr redo command for re-running OCR on
 marked papers

---
 paperforge/cli.py          |   2 +
 paperforge/commands/ocr.py |  62 +++++++++++++
 paperforge/worker/ocr.py   | 185 +++++++++++++++++++++++++++----------
 paperforge/worker/sync.py  |   3 +-
 tests/test_ocr_redo.py     | 177 +++++++++++++++++++++++++++++++++++
 5 files changed, 380 insertions(+), 49 deletions(-)
 create mode 100644 tests/test_ocr_redo.py

diff --git a/paperforge/cli.py b/paperforge/cli.py
index 918ff9c3..f51127bc 100644
--- a/paperforge/cli.py
+++ b/paperforge/cli.py
@@ -231,6 +231,8 @@ def build_parser() -> argparse.ArgumentParser:
     ocr_sub.add_parser("run", help="Run OCR queue")
     doctor_parser = ocr_sub.add_parser("doctor", help="Diagnose OCR configuration and connectivity")
     doctor_parser.add_argument("--live", action="store_true", help="Run live PDF test (L4)")
+    redo_parser = ocr_sub.add_parser("redo", help="Re-run OCR for papers marked ocr_redo: true")
+    redo_parser.set_defaults(ocr_action="redo")
 
     # context (Phase 26: traceable AI context packs)
     p_context = sub.add_parser("context", help="Generate traceable AI context pack for paper(s)")
diff --git a/paperforge/commands/ocr.py b/paperforge/commands/ocr.py
index 0459268a..b325ab3b 100644
--- a/paperforge/commands/ocr.py
+++ b/paperforge/commands/ocr.py
@@ -2,6 +2,8 @@
 
 import argparse
 import logging
+import re as _re
+import shutil
 from pathlib import Path
 
 from paperforge import __version__
@@ -126,6 +128,57 @@ def _get_run_ocr():
     return run_ocr
 
 
+def _run_ocr_redo(vault: Path, verbose: bool = False, no_progress: bool = False) -> int:
+    """Scan for papers with ocr_redo: true, reset their OCR state, then run OCR queue."""
+    from paperforge.adapters.obsidian_frontmatter import extract_preserved_ocr_redo
+    from paperforge.worker._utils import pipeline_paths
+    from paperforge.worker.ocr import run_ocr
+
+    paths = pipeline_paths(vault)
+    ocr_root = paths.get("ocr")
+    lit_root = paths.get("literature")
+
+    if not lit_root or not lit_root.exists():
+        logger.info("No literature directory found, nothing to redo")
+        return 0
+
+    redo_entries = []
+    for note_file in sorted(lit_root.rglob("*.md")):
+        if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
+            continue
+        try:
+            text = note_file.read_text(encoding="utf-8")
+        except Exception:
+            continue
+        if not extract_preserved_ocr_redo(text):
+            continue
+        key_match = _re.search(r"^zotero_key:\s*(.+)$", text, _re.MULTILINE)
+        if not key_match:
+            continue
+        zotero_key = key_match.group(1).strip().strip('"').strip("'")
+        redo_entries.append((zotero_key, note_file, text))
+
+    if not redo_entries:
+        logger.info("No papers with ocr_redo: true found")
+        return 0
+
+    logger.info("Found %d paper(s) with ocr_redo: true", len(redo_entries))
+    for zotero_key, note_file, text in redo_entries:
+        # Delete OCR output directory
+        ocr_dir = ocr_root / zotero_key if ocr_root else None
+        if ocr_dir and ocr_dir.exists():
+            shutil.rmtree(ocr_dir)
+            logger.info("Deleted OCR directory for %s", zotero_key)
+
+        # Update library note frontmatter: ocr_status -> pending, ocr_redo -> false
+        text = _re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=_re.MULTILINE)
+        text = _re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=_re.MULTILINE)
+        note_file.write_text(text, encoding="utf-8")
+        logger.info("Reset ocr_status to pending and ocr_redo to false for %s", zotero_key)
+
+    return run_ocr(vault, verbose=verbose, no_progress=no_progress)
+
+
 def run(args: argparse.Namespace) -> int:
     """Run OCR command.
 
@@ -150,6 +203,15 @@ def run(args: argparse.Namespace) -> int:
     if ocr_action == "doctor" or diagnose_only:
         return _diagnose(vault, live=live, json_output=json_output)
 
+    if ocr_action == "redo":
+        logger.info("OCR redo: scanning for ocr_redo: true papers...")
+        rc = _run_ocr_redo(
+            vault,
+            verbose=getattr(args, "verbose", False),
+            no_progress=getattr(args, "no_progress", False),
+        )
+        return rc
+
     if key:
         logger.info("Processing specific key: %s", key)
 
diff --git a/paperforge/worker/ocr.py b/paperforge/worker/ocr.py
index 76fc92c9..7fb5e553 100644
--- a/paperforge/worker/ocr.py
+++ b/paperforge/worker/ocr.py
@@ -532,6 +532,26 @@ def parse_reference_number(text: str) -> int | None:
     return None
 
 
+def is_reference_lead_block(text: str) -> bool:
+    text = clean_block_text(text)
+    if not text:
+        return False
+    if parse_reference_number(text) is not None:
+        return True
+    if re.match(r"^\s*[^\n]{1,160}\(\d{4}[a-z]?\)\.", text):
+        return True
+    if re.match(r"^\s*[A-Z][A-Za-z'’\-]+(?:\s+[A-Z][A-Za-z'’\-]+)*(?:,\s|\s+et al\.)", text):
+        return True
+    return False
+
+
+def _bbox_column_bucket_xc(block: dict) -> float | None:
+    bbox = block.get("block_bbox", [0, 0, 0, 0])
+    if bbox[2] <= bbox[0]:
+        return None
+    return (bbox[0] + bbox[2]) / 2
+
+
 def sort_reference_blocks(blocks: list[dict]) -> list[dict]:
     numbered_count = sum(parse_reference_number(block.get("block_content", "")) is not None for block in blocks)
     if numbered_count >= max(3, len(blocks) // 2):
@@ -543,9 +563,23 @@ def sort_reference_blocks(blocks: list[dict]) -> list[dict]:
                 block_sort_key(block),
             ),
         )
-    return sorted(
-        blocks, key=lambda block: (block.get("block_bbox", [0, 0, 0, 0])[0], block.get("block_bbox", [0, 0, 0, 0])[1])
-    )
+    xcs = [xc for block in blocks if (xc := _bbox_column_bucket_xc(block)) is not None]
+    if not xcs:
+        mid = 600
+    else:
+        mid = sum(xcs) / len(xcs)
+    left: list[dict] = []
+    right: list[dict] = []
+    for block in blocks:
+        xc = _bbox_column_bucket_xc(block)
+        if xc is None or xc < mid - 60:
+            left.append(block)
+        else:
+            right.append(block)
+    y_sort = lambda block: int(block.get("block_bbox", [0, 0, 0, 0])[1])
+    left.sort(key=y_sort)
+    right.sort(key=y_sort)
+    return left + right
 
 
 def assign_reference_continuation(continuation_block: dict, reference_blocks: list[dict]) -> int | None:
@@ -566,6 +600,53 @@ def assign_reference_continuation(continuation_block: dict, reference_blocks: li
     return candidates[0][0]
 
 
+def build_reference_section_lines(reference_blocks: list[dict], reference_continuations: list[dict]) -> list[str]:
+    if not reference_blocks:
+        return []
+    lines = [""]
+    ordered_reference_blocks = sort_reference_blocks(reference_blocks)
+    continuation_map: dict[int, list[str]] = {}
+    sorted_continuations = sorted(reference_continuations, key=block_sort_key)
+    assigned_continuation_indexes: set[int] = set()
+    incomplete_reference_indexes = [
+        index
+        for index, block in enumerate(ordered_reference_blocks)
+        if clean_block_text(block.get("block_content", "")).rstrip().endswith(":")
+    ]
+    for index, continuation in zip(incomplete_reference_indexes, sorted_continuations, strict=False):
+        continuation_text = clean_block_text(continuation.get("block_content", ""))
+        if continuation_text:
+            continuation_map.setdefault(index, []).append(continuation_text)
+            assigned_continuation_indexes.add(id(continuation))
+    for continuation in sorted_continuations:
+        if id(continuation) in assigned_continuation_indexes:
+            continue
+        target_index = assign_reference_continuation(continuation, ordered_reference_blocks)
+        if target_index is None:
+            continue
+        continuation_map.setdefault(target_index, []).append(clean_block_text(continuation.get("block_content", "")))
+        assigned_continuation_indexes.add(id(continuation))
+    unassigned_continuations = [
+        clean_block_text(continuation.get("block_content", ""))
+        for continuation in sorted_continuations
+        if id(continuation) not in assigned_continuation_indexes
+    ]
+    for index, block in enumerate(ordered_reference_blocks):
+        text = clean_block_text(block.get("block_content", ""))
+        if text:
+            lines.append(text)
+        for continuation_text in continuation_map.get(index, []):
+            if continuation_text:
+                lines.append(continuation_text)
+        if text.rstrip().endswith(":") and unassigned_continuations:
+            lines.append(unassigned_continuations.pop(0))
+    return lines
+
+
+def append_reference_section(rendered: list[str], reference_blocks: list[dict], reference_continuations: list[dict]) -> None:
+    rendered.extend(build_reference_section_lines(reference_blocks, reference_continuations))
+
+
 def clean_author_line(text: str) -> str:
     text = clean_block_text(text)
     text = text.replace("$^{ID}$", "")
@@ -1324,6 +1405,10 @@ def render_page_blocks(
     reference_continuations: list[dict] = []
     footnotes: list[str] = []
     footnote_counter = 0
+    references_heading_seen = False
+    references_section_active = False
+    references_emitted = False
+    references_insert_index: int | None = None
     rendered_cluster_ids: set[int] = set()
     rendered_caption_media_ids: set[int] = set()
     first_page_meta_done = page_index != 1
@@ -1367,7 +1452,16 @@ def render_page_blocks(
                 rendered.extend(deferred_meta)
                 deferred_meta.clear()
                 first_page_meta_done = True
+            if references_heading_seen and (not references_emitted) and title.lower() != "references":
+                if reference_blocks:
+                    append_reference_section(rendered, reference_blocks, reference_continuations)
+                    references_emitted = True
+                references_section_active = False
             rendered.append(f"### {title}")
+            if title.lower() == "references":
+                references_heading_seen = True
+                references_section_active = True
+                references_insert_index = len(rendered)
             continue
         if label == "abstract":
             if page_index == 1 and affiliation_buffer:
@@ -1382,10 +1476,10 @@ def render_page_blocks(
         if label == "reference_content":
             text = clean_block_text(content)
             if text and (not is_reference_tail_noise_line(text)):
-                if parse_reference_number(text) is None:
-                    reference_continuations.append(block)
-                else:
+                if is_reference_lead_block(text):
                     reference_blocks.append(block)
+                else:
+                    reference_continuations.append(block)
             continue
         if label == "text":
             text = clean_block_text(content)
@@ -1397,7 +1491,7 @@ def render_page_blocks(
                 or is_embedded_figure_text_block(block, blocks, page_width=ocr_width, page_height=ocr_height)
             ):
                 continue
-            if raw_reference_blocks and bbox[1] >= first_reference_y - 10:
+            if references_section_active and raw_reference_blocks and bbox[1] >= first_reference_y - 10:
                 reference_continuations.append(block)
                 continue
             if is_numbered_figure_caption(text):
@@ -1525,46 +1619,12 @@ def render_page_blocks(
         rendered.append("")
         rendered.extend(footnotes)
     rendered = dedupe_page_media_lines(rendered)
-    if reference_blocks:
-        rendered.append("")
-        ordered_reference_blocks = sort_reference_blocks(reference_blocks)
-        continuation_map: dict[int, list[str]] = {}
-        sorted_continuations = sorted(reference_continuations, key=block_sort_key)
-        assigned_continuation_indexes: set[int] = set()
-        incomplete_reference_indexes = [
-            index
-            for index, block in enumerate(ordered_reference_blocks)
-            if clean_block_text(block.get("block_content", "")).rstrip().endswith(":")
-        ]
-        for index, continuation in zip(incomplete_reference_indexes, sorted_continuations, strict=False):
-            continuation_text = clean_block_text(continuation.get("block_content", ""))
-            if continuation_text:
-                continuation_map.setdefault(index, []).append(continuation_text)
-                assigned_continuation_indexes.add(id(continuation))
-        for continuation in sorted_continuations:
-            if id(continuation) in assigned_continuation_indexes:
-                continue
-            target_index = assign_reference_continuation(continuation, ordered_reference_blocks)
-            if target_index is None:
-                continue
-            continuation_map.setdefault(target_index, []).append(
-                clean_block_text(continuation.get("block_content", ""))
-            )
-            assigned_continuation_indexes.add(id(continuation))
-        unassigned_continuations = [
-            clean_block_text(continuation.get("block_content", ""))
-            for continuation in sorted_continuations
-            if id(continuation) not in assigned_continuation_indexes
-        ]
-        for index, block in enumerate(ordered_reference_blocks):
-            text = clean_block_text(block.get("block_content", ""))
-            if text:
-                rendered.append(text)
-            for continuation_text in continuation_map.get(index, []):
-                if continuation_text:
-                    rendered.append(continuation_text)
-            if text.rstrip().endswith(":") and unassigned_continuations:
-                rendered.append(unassigned_continuations.pop(0))
+    if reference_blocks and not references_emitted:
+        reference_lines = build_reference_section_lines(reference_blocks, reference_continuations)
+        if references_insert_index is not None:
+            rendered[references_insert_index:references_insert_index] = reference_lines
+        else:
+            rendered.extend(reference_lines)
     return [part for part in rendered if part]
 
 
@@ -1607,7 +1667,7 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu
     return (page_num, markdown_path, json_path, fulltext_md_path)
 
 
-def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False) -> int:
+def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False, redo_mode: bool = False) -> int:
     from paperforge.pdf_resolver import resolve_pdf_path
 
     paths = pipeline_paths(vault)
@@ -1644,6 +1704,35 @@ def run_ocr(vault: Path, verbose: bool = False, no_progress: bool = False) -> in
 
     control_actions = load_control_actions(paths)
     target_keys = {key for key, action in control_actions.items() if action.get("do_ocr", False)}
+
+    if redo_mode:
+        redo_keys = [key for key, action in control_actions.items() if action.get("ocr_redo", False)]
+        if redo_keys:
+            logger.info("OCR redo mode: processing %d paper(s) with ocr_redo: true", len(redo_keys))
+            ocr_root = paths.get("ocr")
+            lit_root = paths.get("literature")
+            for key in redo_keys:
+                ocr_dir = ocr_root / key if ocr_root else None
+                if ocr_dir and ocr_dir.exists():
+                    shutil.rmtree(ocr_dir)
+                    logger.info("Deleted OCR directory for %s", key)
+                target_keys.add(key)
+            if lit_root and lit_root.exists():
+                for note_file in lit_root.rglob("*.md"):
+                    if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
+                        continue
+                    try:
+                        text = note_file.read_text(encoding="utf-8")
+                    except Exception:
+                        continue
+                    for key in redo_keys:
+                        if key in text:
+                            text = re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=re.MULTILINE)
+                            text = re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=re.MULTILINE)
+                            note_file.write_text(text, encoding="utf-8")
+                            logger.info("Reset frontmatter for %s", key)
+                            break
+
     target_rows = []
     for export_path in sorted(paths["exports"].glob("*.json")):
         for item in load_export_rows(export_path):
diff --git a/paperforge/worker/sync.py b/paperforge/worker/sync.py
index 6f825b2d..47c67c2c 100644
--- a/paperforge/worker/sync.py
+++ b/paperforge/worker/sync.py
@@ -198,7 +198,8 @@ def load_control_actions(paths: dict[str, Path]) -> dict[str, dict]:
         )
         if analyze_match:
             analyze = analyze_match.group(1).lower() == "true"
-        actions[zotero_key] = {"analyze": analyze, "do_ocr": do_ocr}
+        ocr_redo = bool(extract_preserved_ocr_redo(text))
+        actions[zotero_key] = {"analyze": analyze, "do_ocr": do_ocr, "ocr_redo": ocr_redo}
     return actions
 
 
diff --git a/tests/test_ocr_redo.py b/tests/test_ocr_redo.py
new file mode 100644
index 00000000..94307943
--- /dev/null
+++ b/tests/test_ocr_redo.py
@@ -0,0 +1,177 @@
+"""Tests for paperforge ocr redo workflow."""
+
+import json
+import re
+import shutil
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _make_vault(tmp_path: Path) -> tuple[Path, Path, Path, Path]:
+    vault = tmp_path / "vault"
+    vault.mkdir()
+    (vault / "paperforge.json").write_text("{}", encoding="utf-8")
+    ocr_root = vault / "PaperForge" / "ocr"
+    ocr_root.mkdir(parents=True)
+    exports = vault / "PaperForge" / "exports"
+    exports.mkdir(parents=True)
+    literature = vault / "Resources" / "Literature"
+    literature.mkdir(parents=True)
+    return vault, ocr_root, exports, literature
+
+
+def _make_ocr_meta(ocr_root: Path, key: str, status: str = "done") -> dict:
+    meta_dir = ocr_root / key
+    meta_dir.mkdir(parents=True, exist_ok=True)
+    (meta_dir / "images").mkdir(exist_ok=True)
+    meta = {
+        "zotero_key": key,
+        "ocr_status": status,
+        "ocr_provider": "PaddleOCR-VL-1.6",
+        "source_pdf": f"some/path/{key}.pdf",
+        "ocr_job_id": "job-123",
+        "ocr_started_at": "2025-01-01T00:00:00",
+        "ocr_finished_at": "2025-01-01T01:00:00",
+        "page_count": 5,
+        "markdown_path": f"PaperForge/ocr/{key}/fulltext.md",
+    }
+    (meta_dir / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
+    return meta
+
+
+def _make_library_note(lit_dir: Path, key: str, ocr_redo: bool = True, ocr_status: str = "done") -> Path:
+    domain_dir = lit_dir / "test_domain"
+    domain_dir.mkdir(parents=True, exist_ok=True)
+    note_path = domain_dir / f"{key}.md"
+    note_text = f"""---
+title: "Test Paper"
+zotero_key: {key}
+do_ocr: true
+analyze: true
+ocr_status: {ocr_status}
+ocr_redo: {"true" if ocr_redo else "false"}
+tags:
+  - test
+---
+
+# Test Paper
+
+Some content
+"""
+    note_path.write_text(note_text, encoding="utf-8")
+    return note_path
+
+
+# ---------------------------------------------------------------------------
+# Test: OCR redo resets ocr_status in meta.json
+# ---------------------------------------------------------------------------
+
+def test_ocr_redo_resets_ocr_status():
+    """Setting ocr_status to pending in meta.json simulates redo reset."""
+    meta = {"zotero_key": "KEY001", "ocr_status": "done", "ocr_job_id": "job-xyz"}
+    meta["ocr_status"] = "pending"
+    meta["ocr_job_id"] = ""
+    assert meta["ocr_status"] == "pending"
+    assert meta["ocr_job_id"] == ""
+    assert meta["zotero_key"] == "KEY001"
+
+
+# ---------------------------------------------------------------------------
+# Test: OCR redo clears OCR directory
+# ---------------------------------------------------------------------------
+
+def test_ocr_redo_clears_ocr_dir(tmp_path):
+    """Verify OCR output directory is removed on redo."""
+    vault, ocr_root, exports, literature = _make_vault(tmp_path)
+    key = "KEY002"
+    _make_ocr_meta(ocr_root, key, status="done")
+    ocr_dir = ocr_root / key
+    assert ocr_dir.exists()
+    assert (ocr_dir / "meta.json").exists()
+    assert (ocr_dir / "images").exists()
+
+    # Simulate redo: delete the OCR directory
+    shutil.rmtree(ocr_dir)
+
+    assert not ocr_dir.exists()
+
+
+# ---------------------------------------------------------------------------
+# Test: OCR redo updates library note frontmatter
+# ---------------------------------------------------------------------------
+
+def test_ocr_redo_updates_frontmatter(tmp_path):
+    """Verify ocr_status and ocr_redo are updated in the library note."""
+    vault, ocr_root, exports, literature = _make_vault(tmp_path)
+    key = "KEY003"
+    _make_ocr_meta(ocr_root, key, status="done")
+    note_path = _make_library_note(literature, key, ocr_redo=True, ocr_status="done")
+
+    # Simulate redo: update frontmatter
+    text = note_path.read_text(encoding="utf-8")
+    text = re.sub(r"^ocr_status:\s*.+$", "ocr_status: pending", text, flags=re.MULTILINE)
+    text = re.sub(r"^ocr_redo:\s*.+$", "ocr_redo: false", text, flags=re.MULTILINE)
+    note_path.write_text(text, encoding="utf-8")
+
+    updated = note_path.read_text(encoding="utf-8")
+    assert "ocr_status: pending" in updated
+    assert "ocr_redo: false" in updated
+    # ocr_redo: false and ocr_redo: true must not both appear
+    assert len(re.findall(r"^ocr_redo:", updated, re.MULTILINE)) == 1
+    assert len(re.findall(r"^ocr_status:", updated, re.MULTILINE)) == 1
+
+
+# ---------------------------------------------------------------------------
+# Test: redo subcommand is registered
+# ---------------------------------------------------------------------------
+
+def test_redo_subcommand_registered():
+    """Verify 'paperforge ocr redo' is a registered subcommand."""
+    from paperforge.cli import build_parser
+
+    parser = build_parser()
+
+    # Parse ocr redo --help should not fail
+    with pytest.raises(SystemExit) as exc:
+        parser.parse_args(["ocr", "redo", "--help"])
+    assert exc.value.code == 0
+
+    # Parse ocr redo should set ocr_action="redo"
+    args = parser.parse_args(["ocr", "redo"])
+    assert args.command == "ocr"
+    assert args.ocr_action == "redo"
+
+
+# ---------------------------------------------------------------------------
+# Test: _run_ocr_redo scan logic
+# ---------------------------------------------------------------------------
+
+def test_ocr_redo_scan_finds_marked_papers(tmp_path):
+    """Verify scan finds papers with ocr_redo: true."""
+    vault, ocr_root, exports, literature = _make_vault(tmp_path)
+    _make_library_note(literature, "KEY_A", ocr_redo=True)
+    _make_library_note(literature, "KEY_B", ocr_redo=False)
+    _make_library_note(literature, "KEY_C", ocr_redo=True)
+
+    from paperforge.adapters.obsidian_frontmatter import extract_preserved_ocr_redo
+
+    found = []
+    for note_file in sorted(literature.rglob("*.md")):
+        if note_file.name in ("fulltext.md", "deep-reading.md", "discussion.md"):
+            continue
+        text = note_file.read_text(encoding="utf-8")
+        if not extract_preserved_ocr_redo(text):
+            continue
+        key_match = re.search(r"^zotero_key:\s*(.+)$", text, re.MULTILINE)
+        assert key_match is not None
+        zkey = key_match.group(1).strip().strip('"').strip("'")
+        found.append(zkey)
+
+    assert "KEY_A" in found
+    assert "KEY_B" not in found
+    assert "KEY_C" in found
+    assert len(found) == 2