Feat/execution backend

backend/api_v2/api_deployment_views.py

@@ -214,9 +214,14 @@ def get(
         response_status = status.HTTP_422_UNPROCESSABLE_ENTITY
         if execution_status_value == CeleryTaskState.COMPLETED.value:
             response_status = status.HTTP_200_OK
-            # Check if highlight data should be removed using configuration registry
+            # Ensure workflow identification keys are always in item metadata
             api_deployment = deployment_execution_dto.api
             organization = api_deployment.organization if api_deployment else None
+            org_id = str(organization.organization_id) if organization else ""
+            DeploymentHelper._enrich_result_with_workflow_metadata(
+                response, organization_id=org_id
+            )
+            # Check if highlight data should be removed using configuration registry
             enable_highlight = False  # Safe default if the key is unavailable (e.g., OSS)
             # Check if the configuration key exists (Cloud deployment) or use settings (OSS)
             from configuration.config_registry import ConfigurationRegistry
@@ -231,8 +236,10 @@ def get(
             if not enable_highlight:
                 response.remove_result_metadata_keys(["highlight_data"])
                 response.remove_result_metadata_keys(["extracted_text"])
-            if not include_metadata:
-                response.remove_result_metadata_keys()
+            if include_metadata or include_metrics:
+                DeploymentHelper._enrich_result_with_usage_metadata(response)
+            if not include_metadata and not include_metrics:
+                response.remove_inner_result_metadata()
             if not include_metrics:
                 response.remove_result_metrics()
         return Response(

backend/api_v2/deployment_helper.py

@@ -258,8 +258,11 @@ def execute_workflow(
             result.status_api = DeploymentHelper.construct_status_endpoint(
                 api_endpoint=api.api_endpoint, execution_id=execution_id
             )
-            # Check if highlight data should be removed using configuration registry
+            # Ensure workflow identification keys are always in item metadata
             organization = api.organization if api else None
+            org_id = str(organization.organization_id) if organization else ""
+            cls._enrich_result_with_workflow_metadata(result, organization_id=org_id)
+            # Check if highlight data should be removed using configuration registry
             enable_highlight = False  # Safe default if the key is unavailable (e.g., OSS)
             from configuration.config_registry import ConfigurationRegistry
 
@@ -273,8 +276,10 @@ def execute_workflow(
             if not enable_highlight:
                 result.remove_result_metadata_keys(["highlight_data"])
                 result.remove_result_metadata_keys(["extracted_text"])
-            if not include_metadata:
-                result.remove_result_metadata_keys()
+            if include_metadata or include_metrics:
+                cls._enrich_result_with_usage_metadata(result)
+            if not include_metadata and not include_metrics:
+                result.remove_inner_result_metadata()
             if not include_metrics:
                 result.remove_result_metrics()
         except Exception as error:
@@ -293,6 +298,120 @@ def execute_workflow(
             )
         return APIExecutionResponseSerializer(result).data
 
+    @staticmethod
+    def _enrich_result_with_usage_metadata(result: ExecutionResponse) -> None:
+        """Enrich each file result's metadata with usage data.
+
+        For each file_execution_id:
+        1. Injects per-model cost arrays (extraction_llm, challenge_llm,
+           embedding) into item["result"]["metadata"].
+        2. Injects aggregated usage totals into item["metadata"]["usage"],
+           matching the legacy response format.
+        """
+        if not isinstance(result.result, list):
+            return
+
+        from usage_v2.helper import UsageHelper
+
+        for item in result.result:
+            if not isinstance(item, dict):
+                continue
+            file_exec_id = item.get("file_execution_id")
+            if not file_exec_id:
+                continue
+
+            # Enrich inner result metadata with per-model breakdown
+            inner_result = item.get("result")
+            if isinstance(inner_result, dict):
+                metadata = inner_result.get("metadata")
+                if isinstance(metadata, dict):
+                    usage_by_model = UsageHelper.get_usage_by_model(file_exec_id)
+                    if usage_by_model:
+                        metadata.update(usage_by_model)
+
+            # Enrich top-level item metadata with aggregated usage
+            item_metadata = item.get("metadata")
+            if isinstance(item_metadata, dict):
+                aggregated = UsageHelper.get_aggregated_token_count(file_exec_id)
+                if aggregated:
+                    aggregated["file_execution_id"] = file_exec_id
+                    item_metadata["usage"] = aggregated
+
+    @staticmethod
+    def _enrich_result_with_workflow_metadata(
+        result: ExecutionResponse,
+        organization_id: str,
+    ) -> None:
+        """Ensure workflow identification keys are always present in item metadata.
+
+        Uses setdefault() — fills in MISSING keys only, never overwrites
+        values already present from the workers cache.
+        """
+        if not isinstance(result.result, list):
+            return
+
+        from workflow_manager.file_execution.models import WorkflowFileExecution
+
+        # 1. Collect file_execution_ids
+        file_exec_ids = [
+            item.get("file_execution_id")
+            for item in result.result
+            if isinstance(item, dict) and item.get("file_execution_id")
+        ]
+        if not file_exec_ids:
+            return
+
+        # 2. Batch query (single JOIN query for all file executions)
+        fe_lookup = {
+            str(fe.id): fe
+            for fe in WorkflowFileExecution.objects.filter(
+                id__in=file_exec_ids
+            ).select_related("workflow_execution")
+        }
+
+        # 3. Get execution-level data (tags) — one M2M query
+        workflow_execution = None
+        tag_names: list[str] = []
+        if fe_lookup:
+            first_fe = next(iter(fe_lookup.values()))
+            workflow_execution = first_fe.workflow_execution
+            tag_names = list(
+                workflow_execution.tags.values_list("name", flat=True)
+            )
+
+        # 4. Enrich each item
+        for item in result.result:
+            if not isinstance(item, dict):
+                continue
+            file_exec_id = item.get("file_execution_id")
+            if not file_exec_id:
+                continue
+
+            # Ensure metadata dict exists
+            if not isinstance(item.get("metadata"), dict):
+                item["metadata"] = {}
+            metadata = item["metadata"]
+
+            fe = fe_lookup.get(str(file_exec_id))
+            we = fe.workflow_execution if fe else workflow_execution
+
+            # Fill MISSING keys only (setdefault won't overwrite)
+            if fe:
+                metadata.setdefault("source_name", fe.file_name)
+                metadata.setdefault("source_hash", fe.file_hash or "")
+                metadata.setdefault("file_execution_id", str(fe.id))
+                metadata.setdefault("total_elapsed_time", fe.execution_time)
+            if we:
+                metadata.setdefault("workflow_id", str(we.workflow_id))
+                metadata.setdefault("execution_id", str(we.id))
+                metadata.setdefault(
+                    "workflow_start_time",
+                    we.created_at.timestamp() if we.created_at else None,
+                )
+
+            metadata.setdefault("organization_id", organization_id)
+            metadata.setdefault("tags", tag_names)
+
     @staticmethod
     def get_execution_status(execution_id: str) -> ExecutionResponse:
         """Current status of api execution.

backend/backend/celery_config.py

@@ -31,3 +31,8 @@ class CeleryConfig:
     beat_scheduler = "django_celery_beat.schedulers:DatabaseScheduler"
 
     task_acks_late = True
+
+    # Prompt Studio IDE callback tasks (ide_index_complete, ide_prompt_complete, etc.)
+    # run on the "prompt_studio_callback" queue, processed by a dedicated Django
+    # backend Celery worker (worker-prompt-studio-callback in docker-compose).
+    # These are sub-second ORM writes + Socket.IO emits after executor completion.

backend/backend/worker_celery.py

@@ -0,0 +1,109 @@
+"""Lightweight Celery app for dispatching tasks to worker-v2 workers.
+
+The Django backend already has a Celery app for internal tasks (beat,
+periodic tasks, etc.) whose broker URL is set via CELERY_BROKER_URL.
+Workers use the same broker. This module provides a second Celery app
+instance that reuses the same broker URL (from Django settings) but
+bypasses Celery's env-var-takes-priority behaviour so it can coexist
+with the main Django Celery app in the same process.
+
+Problem: Celery reads the ``CELERY_BROKER_URL`` environment variable
+with highest priority — overriding constructor args, ``conf.update()``,
+and ``config_from_object()``.
+
+Solution: Subclass Celery and override ``connection_for_write`` /
+``connection_for_read`` so they always use our explicit broker URL,
+bypassing the config resolution chain entirely.
+"""
+
+import logging
+from urllib.parse import quote_plus
+
+from celery import Celery
+from django.conf import settings
+from kombu import Queue
+
+logger = logging.getLogger(__name__)
+
+_worker_app: Celery | None = None
+
+
+class _WorkerDispatchCelery(Celery):
+    """Celery subclass that forces an explicit broker URL.
+
+    Works around Celery's env-var-takes-priority behaviour where
+    ``CELERY_BROKER_URL`` always overrides per-app configuration.
+    The connection methods are the actual points where Celery opens
+    AMQP/Redis connections, so overriding them is both sufficient
+    and safe.
+    """
+
+    _explicit_broker: str | None = None
+
+    def connection_for_write(self, url=None, *args, **kwargs):
+        return super().connection_for_write(
+            url=url or self._explicit_broker, *args, **kwargs
+        )
+
+    def connection_for_read(self, url=None, *args, **kwargs):
+        return super().connection_for_read(
+            url=url or self._explicit_broker, *args, **kwargs
+        )
+
+
+def get_worker_celery_app() -> Celery:
+    """Get or create a Celery app for dispatching to worker-v2 workers.
+
+    The app uses:
+    - Same broker as the workers (built from CELERY_BROKER_BASE_URL,
+      CELERY_BROKER_USER, CELERY_BROKER_PASS via Django settings)
+    - Same PostgreSQL result backend as the Django Celery app
+
+    Returns:
+        Celery app configured for worker-v2 dispatch.
+    """
+    global _worker_app
+    if _worker_app is not None:
+        return _worker_app
+
+    # Reuse the broker URL already built by Django settings (base.py)
+    # from CELERY_BROKER_BASE_URL + CELERY_BROKER_USER + CELERY_BROKER_PASS
+    broker_url = settings.CELERY_BROKER_URL
+
+    # Reuse the same PostgreSQL result backend as Django's Celery app
+    result_backend = (
+        f"db+postgresql://{settings.DB_USER}:"
+        f"{quote_plus(settings.DB_PASSWORD)}"
+        f"@{settings.DB_HOST}:{settings.DB_PORT}/"
+        f"{settings.CELERY_BACKEND_DB_NAME}"
+    )
+
+    app = _WorkerDispatchCelery(
+        "worker-dispatch",
+        set_as_current=False,
+        fixups=[],
+    )
+    # Store the explicit broker URL for use in connection overrides
+    app._explicit_broker = broker_url
+
+    app.conf.update(
+        result_backend=result_backend,
+        task_queues=[Queue("executor")],
+        task_serializer="json",
+        accept_content=["json"],
+        result_serializer="json",
+        result_extended=True,
+    )
+
+    _worker_app = app
+    # Log broker host only (mask credentials)
+    safe_broker = broker_url.split("@")[-1] if "@" in broker_url else broker_url
+    safe_backend = (
+        result_backend.split("@")[-1] if "@" in result_backend else result_backend
+    )
+    logger.info(
+        "Created worker dispatch Celery app (broker=%s, result_backend=%s)",
+        safe_broker,
+        safe_backend,
+    )
+    return _worker_app