Zipstack
diff --git a/‎unstract/core/src/unstract/core/data_models.py‎
Lines changed: 11 additions & 0 deletions b/‎unstract/core/src/unstract/core/data_models.py‎
Lines changed: 11 additions & 0 deletions
diff --git a/‎workers/api-deployment/tasks.py‎
Lines changed: 16 additions & 1 deletion b/‎workers/api-deployment/tasks.py‎
Lines changed: 16 additions & 1 deletion
diff --git a/‎workers/file_processing/tasks.py‎
Lines changed: 58 additions & 19 deletions b/‎workers/file_processing/tasks.py‎
Lines changed: 58 additions & 19 deletions
diff --git a/‎workers/general/tasks.py‎
Lines changed: 10 additions & 0 deletions b/‎workers/general/tasks.py‎
Lines changed: 10 additions & 0 deletions
diff --git a/‎workers/queue_backend/barrier.py‎
Lines changed: 34 additions & 1 deletion b/‎workers/queue_backend/barrier.py‎
Lines changed: 34 additions & 1 deletion
@@ -249,6 +249,17 @@ def normalize_transport(value: object, *, logger: Any = None, context: str = "")
         return DEFAULT_WORKFLOW_TRANSPORT
 
 
+def is_pg_transport(transport: str | None) -> bool:
+    """True if ``transport`` is the Postgres-queue transport.
+
+    Single source for "what counts as PG transport" — centralises the
+    ``== WorkflowTransport.PG_QUEUE.value`` comparison scattered across the
+    worker fan-out / barrier code, and the seam to extend if a second
+    PG-family transport is ever added.
+    """
+    return transport == WorkflowTransport.PG_QUEUE.value
+
+
 class FileListingResult:
     """Result of listing files from a source."""
 
 
@@ -26,7 +26,13 @@
 from shared.workflow.execution.tool_validation import validate_workflow_tool_instances
 from worker import app
 
-from unstract.core.data_models import ExecutionStatus, FileHashData, WorkerFileData
+from unstract.core.data_models import (
+    DEFAULT_WORKFLOW_TRANSPORT,
+    ExecutionStatus,
+    FileHashData,
+    WorkerFileData,
+    normalize_transport,
+)
 from unstract.core.worker_models import ApiDeploymentResultStatus
 
 logger = WorkerLogger.get_logger(__name__)
@@ -683,6 +689,14 @@ def _run_workflow_api(
             org_id=str(schema_name),
             workload_type=WorkloadType.API,
         )
+        # Transport rides in via the dispatched task's kwargs (PR 1 seam); the
+        # fan-out honours it (PG path → fire-and-forget PgBarrier). Fail-closed
+        # to celery on any unrecognised value.
+        transport = normalize_transport(
+            kwargs.get("transport", DEFAULT_WORKFLOW_TRANSPORT),
+            logger=logger,
+            context=f" [exec:{execution_id}]",
+        )
         result = WorkflowOrchestrationUtils.create_chord_execution(
             batch_tasks=batch_tasks,
             callback_task_name="process_batch_callback_api",
@@ -694,6 +708,7 @@ def _run_workflow_api(
             callback_queue=file_processing_callback_queue,
             app_instance=app,
             fairness=api_fairness,
+            transport=transport,
         )
 
         if result is None:
 
@@ -11,6 +11,8 @@
 from typing import Any
 
 from queue_backend import worker_task
+from queue_backend.barrier import BarrierContext
+from queue_backend.pg_barrier import run_batch_with_barrier
 
 # Import shared worker infrastructure
 from shared.api import InternalAPIClient
@@ -222,25 +224,16 @@ def _enhance_batch_with_mrq_flags(
         )
 
 
-def _process_file_batch_core(
-    task_instance, file_batch_data: dict[str, Any]
+def _run_batch_stages(
+    file_batch_data: dict[str, Any], celery_task_id: str
 ) -> dict[str, Any]:
-    """Core implementation of file batch processing.
-
-    This function contains the actual processing logic that both the new task
-    and Django compatibility task will use.
+    """The actual batch work (validate → setup → pre-create → process → compile).
 
-    Args:
-        task_instance: The Celery task instance (self)
-        file_batch_data: Dictionary that will be converted to FileBatchData dataclass
-
-    Returns:
-        Dictionary with successful_files and failed_files counts
+    Transport-agnostic: identical on the Celery chord path and the PG
+    fire-and-forget path. The task instance isn't needed here — its only use
+    (deriving ``celery_task_id``) happens in the caller. Returns the
+    JSON-serialisable batch result.
     """
-    celery_task_id = (
-        task_instance.request.id if hasattr(task_instance, "request") else "unknown"
-    )
-
     # Step 1: Validate and parse input data
     batch_data = _validate_and_parse_batch_data(file_batch_data)
 
@@ -260,6 +253,44 @@ def _process_file_batch_core(
     return _compile_batch_result(context)
 
 
+def _process_file_batch_core(
+    task_instance,
+    file_batch_data: dict[str, Any],
+    barrier_context: BarrierContext | None = None,
+) -> dict[str, Any]:
+    """Core implementation of file batch processing.
+
+    This function contains the actual processing logic that both the new task
+    and Django compatibility task will use.
+
+    Args:
+        task_instance: The Celery task instance (self)
+        file_batch_data: Dictionary that will be converted to FileBatchData dataclass
+        barrier_context: Present only on the 9e PG fire-and-forget path — carries
+            ``execution_id`` / ``batch_index`` / ``callback_descriptor`` so the
+            batch claims its slot and runs the barrier decrement in-body (a
+            PG-consumed task fires no Celery ``.link``). ``None`` on the Celery
+            chord path, where the chord ``.link`` drives the decrement instead.
+
+    Returns:
+        Dictionary with successful_files and failed_files counts
+    """
+    celery_task_id = (
+        task_instance.request.id if hasattr(task_instance, "request") else "unknown"
+    )
+
+    if barrier_context is None:
+        # Celery chord path — the chord's .link runs the decrement after this.
+        return _run_batch_stages(file_batch_data, celery_task_id)
+
+    # PG fire-and-forget path — claim the batch (idempotent on redelivery), run
+    # the stages, then decrement the barrier in-body / self-chain the callback.
+    return run_batch_with_barrier(
+        barrier_context,
+        lambda: _run_batch_stages(file_batch_data, celery_task_id),
+    )
+
+
 @worker_task(
     bind=True,
     name=TaskName.PROCESS_FILE_BATCH,
@@ -272,18 +303,26 @@ def _process_file_batch_core(
     # Timeout inherited from global Celery config (FILE_PROCESSING_TASK_TIME_LIMIT env var)
 )
 @monitor_performance
-def process_file_batch(self, file_batch_data: dict[str, Any]) -> dict[str, Any]:
-    """Process a batch of files in parallel using Celery.
+def process_file_batch(
+    self,
+    file_batch_data: dict[str, Any],
+    _barrier_context: BarrierContext | None = None,
+) -> dict[str, Any]:
+    """Process a batch of files in parallel.
 
     This is the main task entry point for new workers.
 
     Args:
         file_batch_data: Dictionary that will be converted to FileBatchData dataclass
+        _barrier_context: Injected only when this task is dispatched onto the PG
+            queue (9e fire-and-forget path) by ``PgBarrier`` — carries the barrier
+            coordination context (``execution_id`` / ``batch_index`` /
+            ``callback_descriptor``). Absent on the Celery chord path.
 
     Returns:
         Dictionary with successful_files and failed_files counts
     """
-    return _process_file_batch_core(self, file_batch_data)
+    return _process_file_batch_core(self, file_batch_data, _barrier_context)
 
 
 def _validate_and_parse_batch_data(file_batch_data: dict[str, Any]) -> FileBatchData:
 
@@ -56,6 +56,7 @@
     FileBatchData,
     FileHashData,
     WorkerFileData,
+    normalize_transport,
 )
 
 # Import common workflow utilities
@@ -488,6 +489,12 @@ def _execute_general_workflow(
     """
     start_time = time.time()
 
+    # Fail-closed coercion, for parity with the api/scheduler workers (which run
+    # normalize_transport at their entry): a typo'd transport degrades to Celery
+    # with a warning rather than silently routing onto an unknown substrate. The
+    # coerced value feeds both WorkflowContextData and the fan-out below.
+    transport = normalize_transport(transport, logger=logger)
+
     logger.info("Executing general workflow logic for ETL/TASK workflow")
 
     try:
@@ -710,6 +717,7 @@ def _execute_general_workflow(
                 execution_mode=execution_mode,
                 use_file_history=use_file_history,
                 organization_id=api_client.organization_id,
+                transport=transport,
                 **kwargs,
             )
 
@@ -764,6 +772,7 @@ def _orchestrate_file_processing_general(
     execution_mode: tuple | None,
     use_file_history: bool,
     organization_id: str,
+    transport: str = DEFAULT_WORKFLOW_TRANSPORT,
     **kwargs: dict[str, Any],
 ) -> dict[str, Any]:
     """Orchestrate file processing for general workflows using the same pattern as API worker.
@@ -957,6 +966,7 @@ def _orchestrate_file_processing_general(
                 org_id=organization_id,
                 workload_type=WorkloadType.NON_API,
             ),
+            transport=transport,
         )
 
         if not result:
 
@@ -36,10 +36,12 @@
 
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Protocol, TypedDict
+from typing import TYPE_CHECKING, Any, NotRequired, Protocol, TypedDict
 
 from celery import chord
 
+from unstract.core.data_models import DEFAULT_WORKFLOW_TRANSPORT
+
 from .fairness import FairnessKey
 from .handle import BarrierHandle
 
@@ -92,6 +94,29 @@ class CallbackDescriptor(TypedDict):
     kwargs: dict[str, Any]
     queue: str
     fairness_headers: dict[str, Any] | None
+    # 9e: the WorkflowTransport the aggregating callback is fired on when the
+    # barrier completes. Absent / ``None`` → legacy Celery dispatch
+    # (``current_app.apply_async`` — the ``.link`` path). ``"pg_queue"`` → the
+    # fire-and-forget PG path self-chains the callback via dispatch onto PG.
+    # Named ``transport`` (a WorkflowTransport value, e.g. ``"pg_queue"``) — NOT
+    # ``backend``, to avoid confusion with ``QueueBackend`` (``"pg"``).
+    transport: NotRequired[str | None]
+
+
+class BarrierContext(TypedDict):
+    """Per-batch barrier coordination injected into a PG-dispatched header task.
+
+    The 9e fire-and-forget sibling of :class:`CallbackDescriptor`, and typed for
+    the same reason: it crosses producer → PG-queue → consumer, so the contract
+    is pinned here to catch a typo/rename at the type layer rather than as a
+    remote ``KeyError`` mid-batch. Carried as the ``_barrier_context`` kwarg on
+    ``process_file_batch`` so the consumer can claim its slot and run the barrier
+    decrement in-body (a PG-consumed task fires no Celery ``.link``).
+    """
+
+    execution_id: str
+    batch_index: int
+    callback_descriptor: CallbackDescriptor
 
 
 class Barrier(Protocol):
@@ -112,9 +137,15 @@ def enqueue(
         callback_queue: str,
         app_instance: Any,
         fairness: FairnessKey | None = None,
+        transport: str = DEFAULT_WORKFLOW_TRANSPORT,
     ) -> BarrierHandle | None:
         """Enqueue ``header_tasks`` and a single callback to fire on completion.
 
+        ``transport`` is the per-execution transport (9e). Only ``PgBarrier``
+        acts on it (``pg_queue`` → fire-and-forget PG fan-out instead of Celery
+        ``.link``); the Celery/Redis substrates accept it for Protocol parity
+        and ignore it (they are only ever reached on the ``celery`` transport).
+
         ``None`` is the **sole signal** that no work was enqueued — it
         is returned exclusively when ``header_tasks`` is empty. Any
         substrate-level failure (broker outage, serialisation error,
@@ -166,8 +197,10 @@ def enqueue(
         callback_queue: str,
         app_instance: Any,
         fairness: FairnessKey | None = None,
+        transport: str = DEFAULT_WORKFLOW_TRANSPORT,
     ) -> BarrierHandle | None:
         """See :class:`Barrier.enqueue`."""
+        del transport  # Celery chord is only reached on the celery transport.
         # Empty-header guard goes FIRST so a zero-task run skips
         # signature construction / fairness serialisation entirely.
         # Any failure in those paths is now constrained to the