fix: include page citation assets in result zip

apps/worker/tests/contract/test_parse_task_contract.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import json
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from uuid import uuid4
 
+import pandas as pd
 import pytest
 
+from app.services.document_parser.orchestration.parse_output import ParseOutput
 from support.worker_parse_contract import WorkerParseContract
 
 _REPO_ROOT: Path = Path(__file__).resolve().parents[4]
@@ -118,6 +121,108 @@ def test_parse_task_should_process_uploaded_file_through_real_contract_boundarie
     assert contract.find_task_workspaces(tmp_path, job["job_id"]) == []
 
 
+def test_parse_task_result_zip_includes_page_citation_assets(
+    worker_contract_environment: None,
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    contract = WorkerParseContract.create()
+    contract.use_workspace_root(monkeypatch, tmp_path)
+    contract.use_billing(monkeypatch, is_enabled=False)
+
+    source_pdf = tmp_path / "page-citation-source.pdf"
+    _write_blank_pdf(source_pdf, page_count=1)
+
+    job = contract.create_file_job(
+        source_file_name=source_pdf.name,
+        job_id_prefix="job_parse_page_citation",
+    )
+    contract.upload_source_file(
+        local_file_path=source_pdf,
+        s3_key=job["s3_key"],
+    )
+
+    def fake_execute_document_parse(
+        *,
+        job_id: str,
+        job_context: object,
+        prepared_source: object,
+        output_dir: str,
+    ) -> ParseOutput:
+        del job_id, job_context, prepared_source
+        output_path = Path(output_dir)
+        page_citation_assets_dir = output_path / "page_citation_assets"
+        page_citation_assets_dir.mkdir(parents=True, exist_ok=True)
+        (page_citation_assets_dir / "page-1.png").write_bytes(b"page image")
+        (page_citation_assets_dir / "page-999.png").write_bytes(
+            b"unreferenced page image"
+        )
+        parsed_df = pd.DataFrame(
+            [
+                {
+                    "know_id": "page-node-1",
+                    "type": "page",
+                    "content": "Page text",
+                    "path": f"{source_pdf.name}/Root/Introduction",
+                    "length": 9,
+                    "keywords": "",
+                    "summary": "Page summary",
+                    "tokens": "",
+                    "connectto": "",
+                    "page_nums": "1",
+                    "extra_metadata": json.dumps(
+                        {
+                            "page_assets": [
+                                {
+                                    "page_num": 1,
+                                    "artifact_ref": "page_citation_assets/page-1.png",
+                                    "content_type": "image/png",
+                                    "source": "knowhere-rendered-page-citation-source",
+                                    "width": 1200,
+                                    "height": 1800,
+                                }
+                            ],
+                        }
+                    ),
+                }
+            ]
+        )
+        return ParseOutput(output_dir=output_dir, parsed_df=parsed_df)
+
+    monkeypatch.setattr(
+        "app.services.document_ingestion.processing_run.execute_document_parse",
+        fake_execute_document_parse,
+    )
+
+    celery_result = contract.enqueue_parse_task(
+        job_id=job["job_id"],
+        user_id=job["user_id"],
+    )
+
+    assert celery_result.successful()
+    assert celery_result.result["status"] == "success"
+
+    observed = contract.observe_successful_job(job["job_id"])
+    result_row = observed["result"]
+    result_zip = contract.read_result_zip(
+        result_s3_key=result_row["result_s3_key"],
+        tmp_path=tmp_path,
+    )
+
+    assert "page_citation_assets/page-1.png" in result_zip["members"]
+    assert "page_citation_assets/page-999.png" not in result_zip["members"]
+    assert result_zip["chunks"]["chunks"][0]["metadata"]["page_assets"] == [
+        {
+            "page_num": 1,
+            "artifact_ref": "page_citation_assets/page-1.png",
+            "content_type": "image/png",
+            "source": "knowhere-rendered-page-citation-source",
+            "width": 1200,
+            "height": 1800,
+        }
+    ]
+
+
 def test_parse_task_should_charge_user_when_billing_is_enabled(
     worker_contract_environment: None,
     monkeypatch: pytest.MonkeyPatch,

packages/shared-python/shared/services/storage/zip_package_writer.py

@@ -22,6 +22,7 @@ class ZipPackageWriteRequest:
     formatted_chunks: list[dict[str, Any]]
     image_files: tuple[ZipResourceFileInfo, ...]
     table_files: tuple[ZipResourceFileInfo, ...]
+    page_citation_files: tuple[ZipResourceFileInfo, ...]
     doc_nav: dict[str, Any] | None
     manifest: dict[str, Any]
     temp_dir: str | None
@@ -75,6 +76,11 @@ def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact:
 
             self._write_resource_files(zip_file, request.image_files, label="Image")
             self._write_resource_files(zip_file, request.table_files, label="Table")
+            self._write_resource_files(
+                zip_file,
+                request.page_citation_files,
+                label="Page citation asset",
+            )
 
             if request.doc_nav is not None:
                 doc_nav_json = json.dumps(request.doc_nav, ensure_ascii=False, indent=2)

packages/shared-python/shared/services/storage/zip_result_resources.py

@@ -18,6 +18,7 @@
 class ZipPackageResources:
     image_files: tuple[ZipResourceFileInfo, ...]
     table_files: tuple[ZipResourceFileInfo, ...]
+    page_citation_files: tuple[ZipResourceFileInfo, ...]
 
     @property
     def image_files_map(self) -> dict[str, ZipResourceFileInfo]:
@@ -39,9 +40,16 @@ def collect(
     ) -> ZipPackageResources:
         images_dir = os.path.join(add_dir, "images")
         tables_dir = os.path.join(add_dir, "tables")
+        page_citation_assets_dir = os.path.join(add_dir, "page_citation_assets")
         return ZipPackageResources(
             image_files=tuple(self._collect_image_files(chunks, images_dir)),
             table_files=tuple(self._collect_table_files(chunks, tables_dir)),
+            page_citation_files=tuple(
+                self._collect_page_citation_files(
+                    chunks,
+                    page_citation_assets_dir,
+                )
+            ),
         )
 
     def _collect_image_files(
@@ -213,6 +221,51 @@ def _collect_table_files(
 
         return table_files
 
+    def _collect_page_citation_files(
+        self,
+        chunks: list[dict[str, Any]],
+        page_citation_assets_dir: str,
+    ) -> list[ZipResourceFileInfo]:
+        page_citation_refs = _collect_page_citation_refs(chunks)
+        if not page_citation_refs:
+            return []
+        if not os.path.exists(page_citation_assets_dir):
+            raise StorageServiceException(
+                internal_message=(
+                    "Page citation asset directory not found for ZIP packaging: "
+                    f"page_citation_assets_dir={page_citation_assets_dir}"
+                ),
+                operation="collect_page_citation_files",
+            )
+
+        page_citation_files: list[ZipResourceFileInfo] = []
+        for artifact_ref in page_citation_refs:
+            source_path = os.path.join(
+                page_citation_assets_dir,
+                os.path.basename(artifact_ref),
+            )
+            if not os.path.isfile(source_path):
+                raise StorageServiceException(
+                    internal_message=(
+                        "Cannot resolve page citation asset for ZIP packaging: "
+                        f"artifact_ref={artifact_ref}, source_path={source_path}"
+                    ),
+                    operation="collect_page_citation_files",
+                )
+            page_citation_files.append(
+                {
+                    "id": artifact_ref,
+                    "file_path": artifact_ref,
+                    "original_name": os.path.basename(artifact_ref),
+                    "size_bytes": os.path.getsize(source_path),
+                    "format": os.path.splitext(artifact_ref)[1].lstrip("."),
+                    "source_path": source_path,
+                    "zip_path": artifact_ref,
+                }
+            )
+
+        return page_citation_files
+
 def _collect_files_by_name(directory_path: str) -> dict[str, str]:
     files: dict[str, str] = {}
     for filename in os.listdir(directory_path):
@@ -222,6 +275,45 @@ def _collect_files_by_name(directory_path: str) -> dict[str, str]:
     return files
 
 
+def _collect_page_citation_refs(chunks: list[dict[str, Any]]) -> list[str]:
+    refs: list[str] = []
+    seen_refs: set[str] = set()
+    for chunk in chunks:
+        chunk_type = str(chunk.get("type") or "").strip().lower()
+        if chunk_type != "page":
+            continue
+        metadata = chunk.get("metadata")
+        if not isinstance(metadata, dict):
+            continue
+        page_assets = metadata.get("page_assets")
+        if not isinstance(page_assets, list):
+            continue
+        for page_asset in page_assets:
+            if not isinstance(page_asset, dict):
+                continue
+            artifact_ref = _normalize_page_citation_ref(
+                page_asset.get("artifact_ref"),
+            )
+            if artifact_ref and artifact_ref not in seen_refs:
+                seen_refs.add(artifact_ref)
+                refs.append(artifact_ref)
+    return refs
+
+
+def _normalize_page_citation_ref(value: Any) -> str | None:
+    if value is None:
+        return None
+    normalized = str(value).strip().replace("\\", "/").lstrip("/")
+    parts = [
+        part
+        for part in normalized.split("/")
+        if part and part not in {".", ".."}
+    ]
+    if len(parts) != 2 or parts[0] != "page_citation_assets":
+        return None
+    return "/".join(parts)
+
+
 def _add_candidate(candidates: list[str], value: str | None) -> None:
     if not value:
         return

packages/shared-python/shared/services/storage/zip_result_service.py

@@ -87,6 +87,7 @@ def generate_zip_package(
                     formatted_chunks=formatted_chunks,
                     image_files=resources.image_files,
                     table_files=resources.table_files,
+                    page_citation_files=resources.page_citation_files,
                     doc_nav=doc_nav,
                     manifest=manifest,
                     temp_dir=temp_dir,