Fix edge-case bugs found in a thorough code review

app.py

@@ -15,7 +15,7 @@ import asyncio
 from pathlib import Path
 
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Body
-from fastapi.responses import FileResponse, StreamingResponse
+from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 
 from pipeline.jobs import manager
@@ -32,6 +32,31 @@ app = FastAPI(title="Local PDF Text Extractor")
 # Serve static assets (CSS/JS) locally — no external CDN/network dependency.
 app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
 
+# Reject oversized request bodies early — before FastAPI buffers/spools the
+# upload — so a huge (or many concurrent) POST can't exhaust memory. Generous
+# enough to never affect normal local use (the primary, trusted single-user
+# case); it mainly caps the DoS surface of the public demo Space, where body
+# size is attacker-controlled. Checked by Content-Length, so an honest client
+# (browser / curl) is rejected before the body is read.
+MAX_UPLOAD_BYTES = 200 * 1024 * 1024  # 200 MB
+
+
+@app.middleware("http")
+async def _limit_request_body(request, call_next):
+    cl = request.headers.get("content-length")
+    if cl:
+        try:
+            oversized = int(cl) > MAX_UPLOAD_BYTES
+        except ValueError:
+            oversized = False
+        if oversized:
+            return JSONResponse(
+                {"detail": "Request body exceeds the {} MB limit.".format(
+                    MAX_UPLOAD_BYTES // (1024 * 1024))},
+                status_code=413,
+            )
+    return await call_next(request)
+
 # Hold strong references to fire-and-forget background tasks: asyncio keeps only
 # a weak reference to a Task, so a bare create_task() result can in principle be
 # GC'd mid-flight ("Task was destroyed but it is pending"). Keeping them here and
@@ -157,8 +182,15 @@ async def upload(
         )
 
     if handw:
+        # is_available() imports torch/transformers on first call (multi-second
+        # on the CPU build). Offload it to the OCR worker thread — like
+        # capabilities() does — so it never stalls the event loop (which would
+        # freeze other requests' SSE progress). create_job below already awaits
+        # the same executor, so this adds no serialization the upload didn't have.
+        from pipeline.jobs import _EXECUTOR
         from pipeline.handwriting import is_available
-        if not is_available():
+        loop = asyncio.get_running_loop()
+        if not await loop.run_in_executor(_EXECUTOR, is_available):
             raise HTTPException(
                 status_code=400,
                 detail=(

pipeline/analyze.py

@@ -151,18 +151,22 @@ def suggest_settings(pdf_bytes: bytes, *, lang: str = "en",
             return _fallback("PDF has no pages; using defaults.", lang)
 
         # --- Pass 1: classify pages cheaply (text-layer vs needs-OCR). -------
-        # Cap the scan to a small leading window so a 500-page born-digital PDF
-        # (which never accumulates OCR pages) isn't fully scanned — the first
-        # few pages define the document's type for "recommend settings".
+        # Stop as soon as we have enough OCR-eligible pages to SAMPLE for Pass 2.
+        # We deliberately do NOT stop on a leading-window cap: has_text_layer()
+        # is a cheap get_text char-count (no rendering, no OCR), so scanning the
+        # whole doc to confirm it's born-digital is fast — and it avoids the
+        # false "no OCR required" verdict for a typed report with a scanned
+        # appendix / signature page (an OCR page beyond a leading window). A
+        # scanned or mixed doc still short-circuits fast once max_pages OCR pages
+        # are found; only a pure born-digital doc is fully (cheaply) swept.
         text_pages, ocr_page_indices = 0, []
-        scan_cap = max_pages * 4
         for i in range(total):
             page = doc.load_page(i)
             if has_text_layer(page):
                 text_pages += 1
             else:
                 ocr_page_indices.append(i)
-            if len(ocr_page_indices) >= max_pages or i >= min(total - 1, scan_cap):
+            if len(ocr_page_indices) >= max_pages:
                 break
 
         # --- Born-digital short-circuit: no OCR needed. ----------------------

pipeline/export_docx.py

@@ -10,11 +10,14 @@ import math
 import re
 
 # Characters that are illegal in XML 1.0 (and so rejected by python-docx/lxml
-# with "All strings must be XML compatible"). We keep the only valid C0 controls
-# (tab, newline, carriage return) and strip the rest. A form-feed (U+000C) is the
-# common offender: PDF text layers emit it as a page/section separator, so it
-# reaches Word export organically and would otherwise 500 the whole export.
-_XML_ILLEGAL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")
+# with "All strings must be XML compatible", or by the final UTF-8 save with
+# "surrogates not allowed"). We keep the only valid C0 controls (tab, newline,
+# carriage return) and strip the rest. A form-feed (U+000C) is the common
+# offender: PDF text layers emit it as a page/section separator. Lone surrogates
+# (U+D800-U+DFFF) and the non-characters U+FFFE/U+FFFF are the others — PyMuPDF
+# can emit U+FFFE from a broken/byte-swapped ToUnicode CMap, so it reaches Word
+# export organically and would otherwise abort the ENTIRE multi-document export.
+_XML_ILLEGAL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\ud800-\udfff\ufffe\uffff]")
 
 
 def _xml_safe(value) -> str:

pipeline/jobs.py

@@ -441,6 +441,16 @@ class JobManager:
             job.status = "error"
             job.error = str(exc)
             self._broadcast(job, {"type": "error", "error": job.error})
+        finally:
+            # The TTL is an IDLE timer, but last_access is otherwise only set at
+            # creation/stream-start. A job that spent longer than JOB_TTL_SECONDS
+            # actually PROCESSING (e.g. a 1000-page OCR, or a slow online-OCR run)
+            # would otherwise be instantly evictable the moment it finishes — the
+            # next sweep (triggered by any other upload) could drop a result the
+            # user just received, 404-ing a page-image click or KG build. Start
+            # the idle clock at completion instead. Covers every terminal path
+            # (done / cancelled / error / corrupt-open early return).
+            job.last_access = time.monotonic()
 
     def cancel(self, job_id: str) -> None:
         job = self.get(job_id)

pipeline/kg.py

@@ -38,6 +38,7 @@ Public API
 """
 
 import hashlib
+import http.client
 import json
 import re
 import time
@@ -216,6 +217,18 @@ def _post_json(url: str, body: dict, api_key: str, *, timeout: float) -> dict:
             raise RuntimeError(
                 "The Gemini request timed out after {}s.".format(timeout)
             ) from None
+        except (OSError, http.client.HTTPException) as err:
+            # The connection dropped mid-body: urlopen() returned headers but
+            # response.read() failed (e.g. ConnectionResetError — an OSError not
+            # wrapped in URLError — or http.client.IncompleteRead, which isn't
+            # even an OSError). This clause MUST sit after URLError/TimeoutError
+            # (both OSError subclasses). Mapping it to a single-line RuntimeError
+            # keeps the "_post_json only ever raises RuntimeError" contract that
+            # build_graph's per-page isolation and non-fatal embedding rely on.
+            raise RuntimeError(
+                "The Gemini connection dropped while reading the response ({}); "
+                "retry shortly.".format(err)
+            ) from None
     else:  # pragma: no cover - loop always breaks or raises
         raise GeminiHTTPError(
             str(_friendly_http_error(last_http_err)),
@@ -280,6 +293,14 @@ def _list_embedding_models(api_key, *, timeout=30) -> list:
         raise RuntimeError(
             "The Gemini request timed out after {}s.".format(timeout)
         ) from None
+    except (OSError, http.client.HTTPException) as err:
+        # Connection dropped while reading the ListModels body (see _post_json).
+        # Keep it a single-line RuntimeError so _pick_embedding_model's
+        # `except RuntimeError` catches it and falls back to the default model.
+        raise RuntimeError(
+            "The Gemini connection dropped while reading the response ({}); "
+            "retry shortly.".format(err)
+        ) from None
     try:
         payload = json.loads(raw.decode("utf-8", "replace"))
     except (ValueError, AttributeError):
@@ -1006,5 +1027,19 @@ def build_graph(
             graph.set_vectors(vectors, model=resolved)
         except RuntimeError as exc:
             graph.embed_error = str(exc)
+            # If the DISCOVERED model 404'd (Google retired the id mid-process),
+            # drop its per-key cache entry so the NEXT build re-discovers a
+            # working model instead of pinning the dead one for the whole process
+            # lifetime. Only for discovery (no explicit embed_model) and only on a
+            # model-not-found 404 — a transient 429/quota must NOT evict a good id.
+            if (
+                isinstance(exc, GeminiHTTPError)
+                and exc.code == 404
+                and not embed_model
+            ):
+                key_hash = hashlib.sha256(
+                    _clean_key(api_key).encode("utf-8")
+                ).hexdigest()
+                _EMBED_MODEL_RESOLVED.pop(key_hash, None)
 
     return graph

pipeline/online_ocr.py

@@ -185,9 +185,17 @@ def _extract_text(payload: dict) -> str:
     SAFETY / MAX_TOKENS / RECITATION.
     """
     # Prompt blocked before any candidate was produced.
-    feedback = payload.get("promptFeedback") or {}
+    feedback = payload.get("promptFeedback") if isinstance(payload.get("promptFeedback"), dict) else {}
     block_reason = feedback.get("blockReason")
     candidates = payload.get("candidates") or []
+    if not isinstance(candidates, list):
+        # A non-list "candidates" (proxy corruption / future API change) would
+        # otherwise index/attribute-error below and escape the module's
+        # single-line-RuntimeError contract. Degrade to a clean message.
+        raise RuntimeError(
+            "Gemini returned an unexpected response shape (candidates was not a "
+            "list); retry shortly."
+        )
     if not candidates:
         if block_reason:
             raise RuntimeError(
@@ -200,10 +208,14 @@ def _extract_text(payload: dict) -> str:
             "empty; try again or switch models."
         )
 
-    candidate = candidates[0] or {}
+    candidate = candidates[0] if isinstance(candidates[0], dict) else {}
     finish_reason = candidate.get("finishReason")
-    content = candidate.get("content") or {}
-    parts = content.get("parts") or []
+    content = candidate.get("content")
+    if not isinstance(content, dict):
+        content = {}
+    parts = content.get("parts")
+    if not isinstance(parts, list):
+        parts = []
     # ``part.get("text", "")`` only defaults when the key is ABSENT; an explicit
     # ``{"text": null}`` would yield None and break the join. Coerce defensively.
     text = "".join(

pipeline/postprocess.py

@@ -54,6 +54,49 @@ def _key(line: str) -> str:
     return base
 
 
+def _bare_int(line: str):
+    """Return the integer value if ``line`` is a BARE page-number-like integer.
+
+    "Bare" = a single digit group, NO page-word, and nothing but that integer
+    plus surrounding punctuation/whitespace — so "5", "4521", "- 12 -" qualify,
+    but "Page 5" (page-word), "Section 3 intro" (letters) and "1,234.56"
+    (multi-group) do not. Returns ``None`` otherwise. These bare lines are the
+    only ones whose page-number collapse can erase real per-page DATA, so they
+    get the sequentiality guard below; page-word page numbers are always trusted.
+    """
+    base = _WS.sub(" ", line.strip().lower()).strip()
+    if _PAGEWORDS.search(base):
+        return None
+    groups = _DIGITS.findall(base)
+    if len(groups) != 1:
+        return None
+    probe = re.sub(r"[^\w\s]", "", base)
+    if _DIGITS.sub("", probe).strip() != "":  # has letters -> not a bare integer
+        return None
+    try:
+        return int(groups[0])
+    except ValueError:
+        return None
+
+
+def _looks_like_page_numbers(values) -> bool:
+    """Do these per-page integer values look like running page numbers?
+
+    ``values`` is ``None`` for a non-bare-integer key (a page-word page number or
+    a verbatim text header) — those are trusted and kept. For a BARE-integer band
+    the values are treated as page numbers ONLY when they are non-decreasing
+    across pages and span a small range (~ the page count). Arbitrary per-page
+    data (IDs, measurements, amounts) is neither, so it is preserved rather than
+    stripped from every page.
+    """
+    if not values:            # None (not a bare-int key) or empty -> keep as-is
+        return True
+    if len(values) < 2:       # too little signal -> fall back to collapsing
+        return True
+    non_decreasing = all(values[i] <= values[i + 1] for i in range(len(values) - 1))
+    return non_decreasing and (max(values) - min(values)) <= 2 * len(values)
+
+
 def realign_line_confidences(cleaned_text, orig_lines):
     """Map a stripped page's surviving lines back to their confidences.
 
@@ -137,18 +180,34 @@ def strip_running_headers_footers(
         pages_lines.append((lines, nonblank))
 
     top_counter, bot_counter = Counter(), Counter()
+    top_vals, bot_vals = {}, {}  # bare-integer key -> per-page int values (page order)
+
+    def _tally(indices, lines, counter, vals):
+        seen = {}
+        for i in indices:
+            k = _key(lines[i])
+            if k and k not in seen:
+                seen[k] = _bare_int(lines[i])
+        for k, v in seen.items():
+            counter[k] += 1
+            if v is not None:
+                vals.setdefault(k, []).append(v)
+
     for lines, nonblank in pages_lines:
         top, bot = _bands(nonblank, band)
-        for k in {_key(lines[i]) for i in top}:
-            if k:
-                top_counter[k] += 1
-        for k in {_key(lines[i]) for i in bot}:
-            if k:
-                bot_counter[k] += 1
+        _tally(top, lines, top_counter, top_vals)
+        _tally(bot, lines, bot_counter, bot_vals)
 
     thresh = max(2, math.ceil(min_frac * n))
     running_top = {k for k, c in top_counter.items() if c >= thresh}
     running_bot = {k for k, c in bot_counter.items() if c >= thresh}
+    # Value-aware page-number guard: a running key formed from BARE integers is
+    # only treated as a running page number when its per-page values look like
+    # page numbers (see _looks_like_page_numbers). This keeps genuine page
+    # numbers (1,2,3,…) stripped while PRESERVING distinct per-page numeric DATA
+    # (IDs, measurements, amounts) that merely recurs in a top/bottom band.
+    running_top = {k for k in running_top if _looks_like_page_numbers(top_vals.get(k))}
+    running_bot = {k for k in running_bot if _looks_like_page_numbers(bot_vals.get(k))}
     if not running_top and not running_bot:
         cleaned = list(page_texts)
         return (cleaned, _all_indices()) if return_kept else cleaned

static/app.js

@@ -786,6 +786,10 @@
       toolbar.hidden = true;
     }
     refreshExportState();
+    // Removing a card changes the total match set. Re-run the active search so
+    // the "N matches" chip doesn't stay stale (over-counted) — same as the
+    // applyPage and edit-save paths already do.
+    if (jobs.size > 0 && searchInput.value.trim()) runSearch();
   }
 
   function setFileStatus(job, status, label) {