fix(remediate): save pages in batches and resume a killed import

frontend/src/pages/ImportDetail.vue

@@ -127,6 +127,36 @@ function confirmReclassify() {
 	});
 }
 
+// A killed or timed-out job leaves the import stopped with no job behind it.
+const canResume = useCall({
+	url: "/api/v2/method/wikify.api.imports.can_resume_import",
+	method: "GET",
+	immediate: false,
+});
+const mayBeResumable = computed(() => ["Parsing", "Remediating", "Failed"].includes(status.value));
+const isStuck = computed(() => mayBeResumable.value && canResume.data === true);
+watch(
+	status,
+	() => {
+		if (mayBeResumable.value) canResume.submit({ import_name: props.name });
+	},
+	{ immediate: true }
+);
+const resume = useCall({
+	url: "/api/v2/method/wikify.api.imports.resume_import",
+	method: "POST",
+	immediate: false,
+});
+async function resumeImport() {
+	await resume.submit({ import_name: props.name });
+	if (resume.error) {
+		toast.error(resume.error?.messages?.[0] || "Could not resume the import");
+		await canResume.submit({ import_name: props.name });
+	} else {
+		imp.reload();
+	}
+}
+
 // Attach the document (with its project) as the agent's default context. Keyed on the
 // source_document id so reloads (progress ticks) don't reset a page/section selection.
 watch(
@@ -309,8 +339,13 @@ const levelColor = { info: "text-ink-gray-7", warn: "text-ink-amber-6", error: "
 						:route="{ name: 'ImportGraph', params: { name: props.name } }"
 					/>
 					<Dropdown
-						v-if="canRemediate || canReclassify"
+						v-if="canRemediate || canReclassify || isStuck"
 						:options="[
+							{
+								label: 'Resume import',
+								onClick: resumeImport,
+								condition: () => isStuck,
+							},
 							{
 								label: 'Remediate flagged',
 								onClick: () => runRemediation('flagged'),
@@ -333,7 +368,7 @@ const levelColor = { info: "text-ink-gray-7", warn: "text-ink-amber-6", error: "
 							theme="gray"
 							v-bind="actionButtonProps(isMobile, 'lucide-wand-sparkles', 'Actions')"
 							:icon-right="isMobile ? undefined : 'lucide-chevron-down'"
-							:loading="remediate.loading"
+							:loading="remediate.loading || resume.loading"
 						/>
 					</Dropdown>
 					<Dropdown

wikify/api/imports.py

@@ -6,11 +6,13 @@
 
 from wikify.api.sections import assert_not_published
 from wikify.engine import preview_wiki as _preview_wiki
+from wikify.engine import store
 from wikify.jobs import generate as generate_job
 from wikify.jobs._util import IMPORT_JOB_TIMEOUT, import_job_id, log, publish_progress
 from wikify.seed import seed_uncategorized_project
 
 MAX_BATCH = 25
+RESUMABLE_STATUSES = ("Parsing", "Remediating", "Failed")
 
 
 def assert_readable_file(file_url: str) -> None:
@@ -27,14 +29,25 @@ def _create_import(pdf_file_url: str, title: str, project: str) -> str:
 	imp.status = "Queued"
 	imp.insert()
 
+	enqueue_parse(imp.name)
+	return imp.name
+
+
+def enqueue_parse(import_name: str) -> None:
 	frappe.enqueue(
 		"wikify.jobs.parse.run",
 		queue="long",
 		timeout=IMPORT_JOB_TIMEOUT,
-		job_id=import_job_id(imp.name),
-		import_name=imp.name,
+		job_id=import_job_id(import_name),
+		deduplicate=True,
+		import_name=import_name,
 	)
-	return imp.name
+
+
+def is_stuck(import_doc) -> bool:
+	if import_doc.status not in RESUMABLE_STATUSES or is_job_enqueued(import_job_id(import_doc.name)):
+		return False
+	return store.needs_resume(import_doc.source_document)
 
 
 @frappe.whitelist(methods=["POST"])
@@ -87,6 +100,28 @@ def trigger_remediation(import_name: str, scope: str = "flagged") -> str:
 	return import_name
 
 
+@frappe.whitelist()
+def can_resume_import(import_name: str) -> bool:
+	import_doc = frappe.get_doc("Wikify Import", import_name)
+	return import_doc.has_permission("write") and is_stuck(import_doc)
+
+
+@frappe.whitelist(methods=["POST"])
+def resume_import(import_name: str) -> str:
+	import_doc = frappe.get_doc("Wikify Import", import_name)
+	import_doc.check_permission("write")
+	if not is_stuck(import_doc):
+		frappe.throw(
+			_("Only a stopped parse or remediation can be resumed (current status: {0}).").format(
+				import_doc.status
+			)
+		)
+
+	import_doc.db_set({"status": "Queued", "error": None})
+	enqueue_parse(import_name)
+	return import_name
+
+
 @frappe.whitelist(methods=["POST"])
 def reclassify(import_name: str) -> str:
 	imp = frappe.get_doc("Wikify Import", import_name)

wikify/engine/remediate.py

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 
 import fitz
+import frappe
 
 from wikify.engine import diagrams, llm, pdf_utils, regions, settings, store
 from wikify.engine.loader.cleanup_llm import clean_markdown
@@ -16,6 +17,7 @@
 
 ADOPTION_COMPOSITE_RATIO = 0.9
 MIN_CANONICAL_CHARS = 40
+SAVE_BATCH_PAGES = 10
 _IMAGE_EMBED_RE = re.compile(r"(!\[[^\]]*\]\()([^)]*)(\))")
 _MARKUP_RE = re.compile(r"[\s#*_>|`~\-]+")
 
@@ -74,6 +76,7 @@ def remediate_pdf(  # noqa: C901
 	progress_cb: Callable[[int, int], None] | None = None,
 	page_cb: Callable[..., None] | None = None,
 	stage_cb: Callable[[str], None] | None = None,
+	resume: bool = False,
 ) -> dict:
 	if not llm.has_openrouter():
 		raise RuntimeError("OpenRouter key not set — remediation needs cloud models.")
@@ -90,12 +93,21 @@ def remediate_pdf(  # noqa: C901
 	canon_md = {p["page_no"]: p["baseline_markdown"] or "" for p in pages}
 	canon_comp = {p["page_no"]: p["composite"] for p in pages}
 	canon_src = {p["page_no"]: "baseline" for p in pages}
+	saved_page_nos = {p["page_no"] for p in targets if resume and p["canonical_source"]}
+	for p in targets:
+		if p["page_no"] in saved_page_nos:
+			canon_md[p["page_no"]] = p["canonical_markdown"] or ""
+			canon_comp[p["page_no"]] = p["canonical_composite"]
+			canon_src[p["page_no"]] = p["canonical_source"]
+	unsaved_pages: list[dict] = []
 
 	with fitz.open(pdf_path) as doc:
 		furniture = det.find_furniture_lines([doc[p["page_no"] - 1].get_text("text") for p in pages])
 
 		doc_cost = 0.0
 		for i, p in enumerate(targets):
+			if p["page_no"] in saved_page_nos:
+				continue
 			page = doc[p["page_no"] - 1]
 			gt = page.get_text("text")
 			kind = p["kind"]
@@ -193,6 +205,16 @@ def remediate_pdf(  # noqa: C901
 			if progress_cb:
 				progress_cb(i + 1, total)
 
+			unsaved_pages.append(p)
+			if len(unsaved_pages) == SAVE_BATCH_PAGES:
+				with events.suspended_indexing():
+					for saved in unsaved_pages:
+						pno = saved["page_no"]
+						store.set_canonical(saved["name"], canon_md[pno], canon_comp[pno], canon_src[pno])
+				# nosemgrep
+				frappe.db.commit()
+				unsaved_pages = []
+
 	stitched = dict(stitch_cross_page_tables([(p["page_no"], canon_md[p["page_no"]]) for p in pages]))
 	with events.suspended_indexing():
 		for p in pages:

wikify/engine/store.py

@@ -105,7 +105,18 @@ def get_pages(source_document: str) -> list[dict]:
 	return frappe.get_all(
 		"Source Page",
 		filters={"source_document": source_document},
-		fields=["name", "page_no", "kind", "baseline_markdown", "verdict", "composite", "image"],
+		fields=[
+			"name",
+			"page_no",
+			"kind",
+			"baseline_markdown",
+			"verdict",
+			"composite",
+			"image",
+			"canonical_markdown",
+			"canonical_composite",
+			"canonical_source",
+		],
 		order_by="page_no asc",
 	)
 
@@ -140,6 +151,14 @@ def set_canonical(page_name: str, markdown: str, composite: float | None, source
 	invalidate_page(page_name)
 
 
+def needs_resume(source_document: str | None) -> bool:
+	return not source_document or bool(
+		frappe.db.exists(
+			"Source Page", {"source_document": source_document, "canonical_source": ("is", "not set")}
+		)
+	)
+
+
 def set_canonical_mean(source_document: str, mean: float | None) -> None:
 	frappe.db.set_value("Source Document", source_document, "canonical_mean", mean)
 

wikify/jobs/parse.py

@@ -51,7 +51,10 @@ def page_cb(page_no, total, kind, score, metrics) -> None:
 
 		# Parse + score only; the tree is built after remediation (below) so it's built once,
 		# over the cleaned canonical markdown rather than the raw baseline.
-		source_document = parse_pdf(
+		resuming = bool(imp.source_document)
+		if not resuming:
+			delete_partial_documents(import_name)
+		source_document = imp.source_document or parse_pdf(
 			pdf_path,
 			title=imp.import_title,
 			import_name=import_name,
@@ -101,6 +104,7 @@ def rem_page_cb(page_no, total, method, adopted, base_c, new_c, metrics) -> None
 				progress_cb=rem_progress_cb,
 				page_cb=rem_page_cb,
 				stage_cb=stage_cb,
+				resume=resuming,
 			)
 			if result.get("cost"):
 				log(
@@ -139,3 +143,8 @@ def rem_page_cb(page_no, total, method, adopted, base_c, new_c, metrics) -> None
 		frappe.db.commit()
 		log(import_name, "error", "parse", "Parse failed — see error on the import")
 		raise
+
+
+def delete_partial_documents(import_name: str) -> None:
+	for name in frappe.get_all("Source Document", filters={"import": import_name}, pluck="name"):
+		frappe.delete_doc("Source Document", name, ignore_permissions=True, force=True)

wikify/jobs/stalled.py

@@ -4,6 +4,7 @@
 from frappe.utils import add_to_date, now_datetime
 from frappe.utils.background_jobs import is_job_enqueued
 
+from wikify.engine import store
 from wikify.jobs._util import import_job_id, log, publish_progress
 
 RUNNING_STATUSES = ("Queued", "Parsing", "Remediating")
@@ -21,7 +22,11 @@ def fail_stalled_imports() -> None:
 	for row in imports:
 		if is_job_enqueued(import_job_id(row.name)):
 			continue
-		retry_hint = "Click Remediate to retry." if row.source_document else "Import the PDF again to retry."
+		retry_hint = (
+			"Click Resume import to continue."
+			if store.needs_resume(row.source_document)
+			else "Click Remediate to retry."
+		)
 		error = (
 			"The background job stopped before finishing (it timed out or the worker restarted). "
 			+ retry_hint