Refactor worker parse contract boundaries

apps/worker/app/services/document_ingestion/parse_execution.py

@@ -17,7 +17,7 @@ def execute_document_parse(
     prepared_source: PreparedSourceFile,
     output_dir: str,
 ) -> ParseOutput:
-    """Run the parser adapter for a prepared source file."""
+    """Run worker document parsing for a prepared local source file."""
     doc_type = JobMetadataHelper.get_parsing_param(
         job_context.job_metadata,
         "doc_type",

apps/worker/app/services/document_ingestion/processing_billing.py

@@ -24,6 +24,27 @@ class ParseJobBillingSnapshot:
     billing_status: str
 
 
+def record_skipped_parse_job_billing(
+    *,
+    job_id: str,
+    workload_estimate: WorkloadEstimate,
+) -> ParseJobBillingSnapshot:
+    page_count = workload_estimate.page_count
+    with get_sync_db_context() as db:
+        job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update())
+        job = job_result.scalar_one_or_none()
+        if job:
+            job.page_count = page_count
+            job.credits_charged = 0
+            job.billing_status = "skipped"
+
+    return ParseJobBillingSnapshot(
+        billing_amount_micro_dollars=0,
+        billing_credits=0.0,
+        billing_status="skipped",
+    )
+
+
 def charge_parse_job_pages(
     *,
     job_id: str,
@@ -121,5 +142,13 @@ def record_processing_start(
         metadata_updates["workload_estimate_fallback_reason"] = (
             workload_estimate.fallback_reason
         )
+    with get_sync_db_context() as db:
+        job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update())
+        job = job_result.scalar_one_or_none()
+        if job:
+            job.job_metadata = {
+                **dict(job.job_metadata or {}),
+                **metadata_updates,
+            }
     job_context.metadata_service.update_metadata(job_id, metadata_updates)
     job_context.job_metadata.update(metadata_updates)

apps/worker/app/services/document_ingestion/processing_context.py

@@ -1,24 +1,21 @@
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from typing import Any
 
 from loguru import logger
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 
-from shared.core.config import settings
 from shared.core.database_sync import get_sync_db_context
 from shared.core.exceptions.domain_exceptions import (
     NotFoundException,
-    ValidationException,
 )
 from shared.models.database.job import Job
 from shared.services.redis.redis_sync_service import (
     SyncJobInfoRedisService,
     SyncJobMetadataService,
 )
-from shared.services.storage.job_file_storage import JobFileStorage
 
 
 @dataclass(frozen=True)
@@ -37,15 +34,13 @@ def load_parse_job_context(
 ) -> ParseJobContext:
     job_info_service = SyncJobInfoRedisService(redis_service)
     job_info = job_info_service.get_job_info(job_id)
+    job_row: Job | None = None
 
     if not job_info:
         logger.warning(
             f"JobInfo not found in Redis for job_id={job_id}; falling back to database"
         )
-        with get_sync_db_context() as fallback_db:
-            job_row = fallback_db.execute(
-                select(Job).where(Job.job_id == job_id)
-            ).scalar_one_or_none()
+        job_row = _load_job_row(job_id)
 
         if not job_row or not job_row.s3_key:
             raise NotFoundException(
@@ -58,6 +53,17 @@ def load_parse_job_context(
         job_user_id: str | None = (
             str(job_row.user_id) if job_row.user_id else requested_user_id
         )
+        job_info_service.save_job_info(
+            job_id,
+            {
+                "job_id": job_id,
+                "s3_key": s3_key,
+                "user_id": job_user_id,
+                "webhook_enabled": bool(job_row.webhook_enabled),
+                "job_type": "document_ingestion",
+                "source_type": job_row.source_type,
+            },
+        )
         logger.info(f"Recovered JobInfo from database: job_id={job_id}, s3_key={s3_key}")
     else:
         raw_s3_key = job_info.get("s3_key")
@@ -77,11 +83,17 @@ def load_parse_job_context(
     metadata_service = SyncJobMetadataService(redis_service)
     raw_job_metadata = metadata_service.get_metadata(job_id)
     if not isinstance(raw_job_metadata, dict) or not raw_job_metadata:
-        raise NotFoundException(
-            resource="JobMetadata",
-            resource_id=job_id,
-            internal_message=f"Job metadata not found for job_id={job_id}",
-        )
+        job_row = job_row or _load_job_row(job_id)
+        raw_job_metadata = job_row.job_metadata if job_row else None
+        if isinstance(raw_job_metadata, dict) and raw_job_metadata:
+            metadata_service.save_metadata(job_id, raw_job_metadata)
+            logger.info(f"Recovered JobMetadata from database: job_id={job_id}")
+        else:
+            raise NotFoundException(
+                resource="JobMetadata",
+                resource_id=job_id,
+                internal_message=f"Job metadata not found for job_id={job_id}",
+            )
 
     return ParseJobContext(
         job_metadata=dict(raw_job_metadata),
@@ -92,30 +104,10 @@ def load_parse_job_context(
     )
 
 
-def assert_source_file_within_size_limit(s3_key: str) -> None:
-    file_info = JobFileStorage().verify_upload_exists(s3_key)
-    if not file_info.get("exists"):
-        raise NotFoundException(
-            resource="S3File",
-            resource_id=s3_key,
-            internal_message=f"S3 file not found: {s3_key}",
-        )
+def _load_job_row(job_id: str) -> Job | None:
+    with get_sync_db_context() as fallback_db:
+        return _select_job_row(fallback_db, job_id)
 
-    logger.info(f"S3 file verified: {s3_key}")
-
-    file_size = file_info.get("size", 0)
-    file_extension = os.path.splitext(s3_key)[1].lower()
-    if file_size > settings.MAX_FILE_SIZE:
-        limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024)
-        raise ValidationException(
-            user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})",
-            violations=[
-                {
-                    "field": "file_size",
-                    "description": (
-                        f"Size {file_size} bytes exceeds limit of "
-                        f"{settings.MAX_FILE_SIZE} bytes"
-                    ),
-                }
-            ],
-        )
+
+def _select_job_row(db: Session, job_id: str) -> Job | None:
+    return db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none()

apps/worker/app/services/document_ingestion/processing_run.py

@@ -11,22 +11,21 @@
 from app.services.document_ingestion.parse_execution import execute_document_parse
 from app.services.document_ingestion.processing_billing import (
     charge_parse_job_pages,
+    record_skipped_parse_job_billing,
     record_processing_start,
 )
 from app.services.document_ingestion.processing_context import (
     ParseJobContext,
-    assert_source_file_within_size_limit,
     load_parse_job_context,
 )
 from app.services.document_ingestion.source_preparation import prepare_source_file
 from app.services.document_ingestion.success_finalization import finalize_parse_success
 from app.services.document_ingestion.workspace import (
     TemporaryParseWorkspace,
-    cleanup_task_workspace,
-    download_s3_file_to_temp,
 )
 from loguru import logger
 
+from shared.core.exceptions.domain_exceptions import ValidationException
 from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service
 from shared.services.redis.distributed_lock import RedisJobLock
 from shared.services.redis.redis_sync_service import (
@@ -44,7 +43,6 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]:
 
         redis_service = SyncRedisServiceFactory.get_service()
         job_context = load_parse_job_context(job_id, user_id, redis_service)
-        assert_source_file_within_size_limit(job_context.s3_key)
 
         should_process = mark_job_running(job_id, job_context.redis_service)
         if not should_process:
@@ -65,7 +63,7 @@ def execute(self, job_id: str, user_id: str | None) -> dict[str, object]:
                     task_workspace=task_workspace,
                 )
             finally:
-                task_workspace.cleanup(cleanup_task_workspace)
+                task_workspace.cleanup()
 
         return result
 
@@ -83,7 +81,6 @@ def _run_parse_job(
         job_id=job_id,
         job_context=job_context,
         input_dir=task_workspace.input_dir,
-        download_source_file=download_s3_file_to_temp,
     )
 
     workload_estimate = PageEstimator.estimate_workload(prepared_source.local_file_path)
@@ -96,6 +93,23 @@ def _run_parse_job(
     )
 
     processing_started_at = datetime.now(timezone.utc)
+    if _is_pdf_page_limit_exceeded(
+        file_extension=prepared_source.file_extension,
+        page_count=page_count,
+    ):
+        billing_snapshot = record_skipped_parse_job_billing(
+            job_id=job_id,
+            workload_estimate=workload_estimate,
+        )
+        record_processing_start(
+            job_id=job_id,
+            job_context=job_context,
+            billing_snapshot=billing_snapshot,
+            processing_started_at=processing_started_at,
+            workload_estimate=workload_estimate,
+        )
+        _raise_pdf_page_limit_exceeded(page_count)
+
     billing_snapshot = charge_parse_job_pages(
         job_id=job_id,
         filename=prepared_source.source_file_name,
@@ -146,3 +160,27 @@ def _run_parse_job(
         task_workspace_dir=task_workspace.root_dir,
         result_storage_factory=get_result_storage,
     )
+
+
+def _is_pdf_page_limit_exceeded(*, file_extension: str, page_count: int) -> bool:
+    from shared.core.config import settings
+
+    return file_extension == ".pdf" and page_count > settings.MAX_PDF_PAGE_LIMIT
+
+
+def _raise_pdf_page_limit_exceeded(page_count: int) -> None:
+    from shared.core.config import settings
+
+    pdf_page_limit = settings.MAX_PDF_PAGE_LIMIT
+    raise ValidationException(
+        user_message=(
+            f"Document too large: {page_count} pages exceeds the {pdf_page_limit}-page limit. "
+            "Please split the document and upload in smaller batches."
+        ),
+        violations=[
+            {
+                "field": "page_count",
+                "description": f"PDF has {page_count} pages, limit is {pdf_page_limit}",
+            }
+        ],
+    )

apps/worker/app/services/document_ingestion/source_preparation.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 
 import os
-from collections.abc import Callable
 from dataclasses import dataclass
 
 from app.services.document_ingestion.processing_context import ParseJobContext
-from app.services.document_ingestion.workspace import download_s3_file_to_temp
 from app.services.document_parser.support.internal_parse_name import (
     prepare_internal_parse_input,
 )
 from loguru import logger
 
+from shared.core.config import settings
+from shared.core.exceptions.domain_exceptions import (
+    NotFoundException,
+    ValidationException,
+)
 from shared.models.schemas.job_metadata import JobMetadataHelper
-
-DownloadSourceFile = Callable[[str, str, str], str]
+from shared.services.storage.job_file_storage import JobFileStorage
 
 
 @dataclass(frozen=True)
@@ -29,20 +31,21 @@ def prepare_source_file(
     job_id: str,
     job_context: ParseJobContext,
     input_dir: str,
-    download_source_file: DownloadSourceFile = download_s3_file_to_temp,
 ) -> PreparedSourceFile:
-    """Download and normalize the source file before parser execution."""
+    """Verify, download, and normalize the uploaded source file."""
     source_file_name = JobMetadataHelper.get_source_file_name(
         job_context.job_metadata,
     ) or os.path.basename(job_context.s3_key)
     file_extension = (
         os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else ""
     )
 
-    local_file_path = download_source_file(
+    storage = JobFileStorage()
+    _assert_source_file_within_size_limit(storage, job_context.s3_key)
+    local_file_path = storage.download_upload_to_temp(
         job_context.s3_key,
-        file_extension,
-        input_dir,
+        suffix=file_extension,
+        temp_dir=input_dir,
     )
     logger.info(f"File downloaded: job_id={job_id}, local_path={local_file_path}")
 
@@ -64,3 +67,36 @@ def prepare_source_file(
         local_file_path=prepared_parse_input.file_path,
         file_extension=file_extension,
     )
+
+
+def _assert_source_file_within_size_limit(
+    storage: JobFileStorage,
+    s3_key: str,
+) -> None:
+    file_info = storage.verify_upload_exists(s3_key)
+    if not file_info.get("exists"):
+        raise NotFoundException(
+            resource="S3File",
+            resource_id=s3_key,
+            internal_message=f"S3 file not found: {s3_key}",
+        )
+
+    logger.info(f"S3 file verified: {s3_key}")
+
+    raw_file_size = file_info.get("size", 0)
+    file_size = raw_file_size if isinstance(raw_file_size, int) else 0
+    file_extension = os.path.splitext(s3_key)[1].lower()
+    if file_size > settings.MAX_FILE_SIZE:
+        limit_mb = settings.MAX_FILE_SIZE // (1024 * 1024)
+        raise ValidationException(
+            user_message=f"File size exceeds limit (max {limit_mb}MB for {file_extension})",
+            violations=[
+                {
+                    "field": "file_size",
+                    "description": (
+                        f"Size {file_size} bytes exceeds limit of "
+                        f"{settings.MAX_FILE_SIZE} bytes"
+                    ),
+                }
+            ],
+        )