fix: harden background workflow result handling

Author: erichare

src/backend/base/langflow/api/v2/workflow.py

@@ -657,40 +657,55 @@ async def get_workflow_status(
             persisted_request = (job.job_metadata or {}).get("request") or {}
             effective_session_id = persisted_request.get("session_id") or flow_id_str
 
-            # Default GET-status path: rebuild from the durable ``output`` events
-            # the runner captured into ``Job.result`` (langflow-protocol). This
-            # needs no ``vertex_build`` rows, so it works with vertex-build storage
-            # off (headless), and skips the graph reconstruction the vertex-build
-            # path does.
+            # Default GET-status path: rebuild from the protocol-neutral terminal
+            # captures the runner stored in ``Job.result``. This needs no
+            # ``vertex_build`` rows, so it works with vertex-build storage off
+            # (headless) and skips graph reconstruction.
             result = job.result if isinstance(job.result, dict) else {}
             output_events = result.get("outputs") or []
-            if output_events:
-                return workflow_response_from_output_events(
-                    output_events,
-                    flow_id=flow_id_str,
-                    job_id=job_id_str,
-                    session_id=effective_session_id,
-                )
+            partial_stored_response: WorkflowExecutionResponse | None = None
+            if isinstance(output_events, list) and output_events:
+                try:
+                    return workflow_response_from_output_events(
+                        output_events,
+                        flow_id=flow_id_str,
+                        job_id=job_id_str,
+                        session_id=effective_session_id,
+                        fail_on_rejected=True,
+                    )
+                except ValueError:
+                    # Preserve every decodable capture in case the legacy
+                    # vertex-build fallback is also unavailable.
+                    partial_stored_response = workflow_response_from_output_events(
+                        output_events,
+                        flow_id=flow_id_str,
+                        job_id=job_id_str,
+                        session_id=effective_session_id,
+                    )
 
-            # Fallback: ``Job.result`` carried no outputs (an agui-protocol run
-            # leaves it empty; the terminal output lives only on /events) or a
-            # legacy job predates output capture. Reconstruct from the
-            # ``vertex_build`` rows keyed by job_id when they exist; if none do,
-            # ValueError degrades to a bare COMPLETED with an empty outputs map.
+            # Fallback: ``Job.result`` carried no outputs, has an invalid shape,
+            # contains rejected/version-skewed entries, or predates capture.
+            # Reconstruct from ``vertex_build`` rows keyed by job_id when present.
             try:
-                return await reconstruct_workflow_response_from_job_id(
+                reconstructed = await reconstruct_workflow_response_from_job_id(
                     session=session,
                     flow=flow,
                     job_id=job_id_str,
                     user_id=str(current_user.id),
                 )
             except ValueError:
+                if partial_stored_response is not None:
+                    return partial_stored_response
                 return workflow_response_from_output_events(
                     [],
                     flow_id=flow_id_str,
                     job_id=job_id_str,
                     session_id=effective_session_id,
                 )
+            else:
+                if reconstructed.session_id is None:
+                    reconstructed = reconstructed.model_copy(update={"session_id": effective_session_id})
+                return reconstructed
 
         if job.status == JobStatus.FAILED:
             # Surface the durable error JSON the runner persisted, additively.

src/backend/base/langflow/services/background_execution/runner.py

@@ -147,7 +147,6 @@ async def _suspend(self, job_id: UUID, exc: PauseRequested) -> None:
         finalizer, or close the live bus — the run is resumable.
         """
         await self._jobs.append_event(job_id, HUMAN_INPUT_REQUIRED_EVENT, exc.payload)
-        await self._jobs.update_job_status(job_id, JobStatus.SUSPENDED)
         metadata: dict[str, Any] = {}
         if exc.request_id is not None:
             metadata["pending_request_id"] = exc.request_id
@@ -159,8 +158,10 @@ async def _suspend(self, job_id: UUID, exc: PauseRequested) -> None:
             # pre-seed its capture list. Deduped so a resume-pause-resume cycle does
             # not accumulate duplicates (the pre-seed already folds the prior stash in).
             metadata["pre_pause_outputs"] = self._dedup_outputs(exc.output_events)
-        if metadata:
-            await self._jobs.update_job_metadata(job_id, metadata)
+        # Publish SUSPENDED only when the resume guard and pre-pause captures are
+        # visible in the same committed row. Otherwise a racing resume can claim
+        # the job between separate status and metadata transactions.
+        await self._jobs.suspend_job(job_id, metadata)
 
     async def _finalize_terminal_event(self, job_id: UUID) -> None:
         """Backfill the error blob + terminal event for TIMED_OUT / CANCELLED.
@@ -238,12 +239,11 @@ async def _drive(self, *, job_id: UUID, source_kwargs: dict[str, Any]) -> None:
         last_durable_seq = 0
         errored_payload: dict[str, Any] | None = None
         # Terminal outputs the run produced, captured so GET status can return
-        # the result without forcing a /events re-attach. The langflow adapter
-        # normalizes each terminal output into a durable ``output`` event whose
-        # ``data`` is the same ``OutputEvent`` (a ``ComponentOutput`` plus its
-        # component id) that sync returns in ``outputs[id]``. The agui adapter
-        # does not emit these, so agui-protocol runs leave this empty and their
-        # status stays result-less (the result is still on the /events log).
+        # the result without forcing a /events re-attach. The frame source emits
+        # a protocol-neutral off-wire capture for both langflow and AG-UI. The
+        # public langflow ``output`` event is also accepted as a fallback for
+        # injected/exported FrameSource implementations that do not emit the
+        # private capture; final deduplication keeps the two paths from doubling.
         #
         # A resumed run starts a FRESH frame stream and the output loop skips
         # re-emitting checkpoint-restored vertices, so any terminal output the
@@ -256,7 +256,7 @@ async def _drive(self, *, job_id: UUID, source_kwargs: dict[str, Any]) -> None:
                 # Off-wire terminal-output capture (both protocols): record it into
                 # the in-memory result ONLY. Deliberately no ``append_event`` (never
                 # in ``job_events``) and no ``publish`` (never on the live bus), so the
-                # wire is unchanged and ``Job.result`` fills from these captures alone.
+                # wire is unchanged and both protocols populate ``Job.result``.
                 payload = self._decode_payload(frame_bytes)
                 output_data = payload.get("data")
                 if isinstance(output_data, dict):
@@ -287,13 +287,14 @@ async def _drive(self, *, job_id: UUID, source_kwargs: dict[str, Any]) -> None:
                 seq = await self._jobs.append_event(job_id, event_type, payload)
                 last_durable_seq = seq
                 await self._bus.publish(str(job_id), LiveFrame(seq=seq, data=self._restamp_id(frame_bytes, seq)))
+                if event_type == "output":
+                    # Fallback for custom frame sources that emit the standard
+                    # adapter event but not WORKFLOW_OUTPUT_CAPTURE_EVENT.
+                    output_data = payload.get("data")
+                    if isinstance(output_data, dict):
+                        output_events.append(output_data)
                 if event_type == self._adapter.terminal_error_type:
                     errored_payload = payload
-                # NOTE: terminal outputs are captured off-wire via the
-                # WORKFLOW_OUTPUT_CAPTURE_EVENT frame above (protocol-neutral), not
-                # from the langflow adapter's durable wire ``output`` event — that
-                # left agui-protocol ``Job.result`` empty. The wire ``output`` still
-                # flows to streaming clients via ``append_event``/``publish`` here.
             else:
                 await self._bus.publish(
                     str(job_id),

src/backend/base/langflow/services/jobs/service.py

@@ -248,6 +248,26 @@ async def update_job_metadata(
             await session.flush()
             return job
 
+    async def suspend_job(self, job_id: UUID, metadata: dict) -> Job | None:
+        """Atomically merge pause metadata and transition a job to ``SUSPENDED``.
+
+        Resume-critical metadata must become visible in the same transaction as
+        the status transition. Otherwise a status reader can observe
+        ``SUSPENDED`` before ``pending_request_id`` or ``pre_pause_outputs`` is
+        committed and return an incomplete pause response.
+
+        Returns the updated Job, or ``None`` if the row does not exist.
+        """
+        async with session_scope() as session:
+            job = await session.get(Job, job_id)
+            if job is None:
+                return None
+            job.job_metadata = {**(job.job_metadata or {}), **metadata}
+            job.status = JobStatus.SUSPENDED
+            session.add(job)
+            await session.flush()
+            return job
+
     async def set_result(self, job_id: UUID, result: dict | None) -> Job | None:
         """Persist the durable terminal result blob for a job.
 

src/backend/tests/unit/api/v2/test_workflow.py

@@ -19,7 +19,7 @@
 from httpx import AsyncClient
 from langflow.services.database.models.flow.model import Flow
 from langflow.services.database.models.jobs.model import Job, JobType
-from lfx.schema.workflow import JobStatus
+from lfx.schema.workflow import JobStatus, WorkflowExecutionResponse
 from lfx.services.deps import session_scope
 from sqlalchemy.exc import OperationalError
 
@@ -203,7 +203,7 @@ async def test_get_status_completed_reconstruction(
         client: AsyncClient,
         created_api_key,
     ):
-        """Test GET /workflow returns reconstructed response for a completed job."""
+        """Vertex-build fallback keeps the submitted session for a data-only result."""
         job_id = uuid4()
 
         flow_id = uuid4()
@@ -213,6 +213,8 @@ async def test_get_status_completed_reconstruction(
         mock_job.status = JobStatus.COMPLETED
         mock_job.type = JobType.WORKFLOW
         mock_job.user_id = None
+        mock_job.result = None
+        mock_job.job_metadata = {"request": {"session_id": "submitted-session"}}
 
         with (
             patch("langflow.api.v2.workflow.get_job_service") as mock_get_job_service,
@@ -227,16 +229,86 @@ async def test_get_status_completed_reconstruction(
             mock_flow.id = flow_id
             mock_get_flow.return_value = mock_flow
 
-            mock_reconstruct.return_value = {"flow_id": str(flow_id), "status": "completed", "outputs": {}}
+            mock_reconstruct.return_value = WorkflowExecutionResponse(
+                flow_id=str(flow_id),
+                status=JobStatus.COMPLETED,
+                outputs={},
+                session_id=None,
+            )
 
             headers = {"x-api-key": created_api_key.api_key}
             response = await client.get(f"api/v2/workflows?job_id={job_id}", headers=headers)
 
             assert response.status_code == 200
             result = response.json()
             assert result["status"] == "completed"
+            assert result["session_id"] == "submitted-session"
             mock_reconstruct.assert_called_once()
 
+    async def test_get_status_malformed_stored_output_falls_back_to_vertex_builds(
+        self,
+        client: AsyncClient,
+        created_api_key,
+    ):
+        """A truthy but incomplete Job.result must not suppress valid vertex-build recovery."""
+        job_id = uuid4()
+        flow_id = uuid4()
+        mock_job = MagicMock(
+            job_id=job_id,
+            flow_id=flow_id,
+            status=JobStatus.COMPLETED,
+            type=JobType.WORKFLOW,
+            user_id=None,
+            result={
+                "outputs": [
+                    {"component_id": 123, "type": "message", "status": "completed", "content": "ignored"},
+                    {"component_id": [], "type": "message", "status": "completed", "content": "ignored"},
+                ]
+            },
+            job_metadata={"request": {"session_id": "submitted-session"}},
+        )
+
+        with (
+            patch("langflow.api.v2.workflow.get_job_service") as mock_get_job_service,
+            patch("langflow.api.v2.workflow.get_flow_by_id_or_endpoint_name") as mock_get_flow,
+            patch("langflow.api.v2.workflow.reconstruct_workflow_response_from_job_id") as mock_reconstruct,
+        ):
+            mock_service = MagicMock()
+            mock_service.get_job_by_job_id = AsyncMock(return_value=mock_job)
+            mock_get_job_service.return_value = mock_service
+
+            mock_flow = MagicMock(id=flow_id)
+            mock_get_flow.return_value = mock_flow
+            mock_reconstruct.return_value = WorkflowExecutionResponse(
+                flow_id=str(flow_id),
+                status=JobStatus.COMPLETED,
+                outputs={
+                    "DataOutput-1": {
+                        "type": "data",
+                        "status": JobStatus.COMPLETED,
+                        "content": {"ok": True},
+                    }
+                },
+            )
+
+            response = await client.get(
+                f"api/v2/workflows?job_id={job_id}",
+                headers={"x-api-key": created_api_key.api_key},
+            )
+
+        assert response.status_code == 200
+        assert response.json()["outputs"] == {
+            "DataOutput-1": {
+                "type": "data",
+                "status": "completed",
+                "display_name": None,
+                "content": {"ok": True},
+                "metadata": None,
+            }
+        }
+        assert response.json()["session_id"] == "submitted-session"
+        mock_reconstruct.assert_awaited_once()
+
     async def test_get_status_timed_out(
         self,
         client: AsyncClient,

src/backend/tests/unit/api/v2/test_workflow_background.py

@@ -132,7 +132,7 @@ async def test_background_status_returns_output(client, created_api_key, bg_flow
     assert body["outputs"], f"completed background status carried no outputs: {body}"
 
 
-async def test_background_status_from_job_table_with_vertex_builds_off(client, created_api_key, bg_flow):
+async def test_background_status_from_job_table_with_vertex_builds_off(client, created_api_key, bg_flow, monkeypatch):
     """With vertex_build storage OFF, GET status must still carry the full output.
 
     Proves the headless-executor path: disable ``vertex_builds_storage_enabled`` so
@@ -150,49 +150,45 @@ async def test_background_status_from_job_table_with_vertex_builds_off(client, c
     from lfx.services.deps import get_settings_service
 
     settings = get_settings_service().settings
-    original = settings.vertex_builds_storage_enabled
-    settings.vertex_builds_storage_enabled = False
-    try:
-        submit = await client.post("api/v2/workflows", json=_body(bg_flow), headers=_headers(created_api_key))
-        assert submit.status_code == 200, submit.text
-        job_id = submit.json()["job_id"]
-
-        row = None
-        for _ in range(200):
-            async with session_scope() as session:
-                row = await session.get(Job, UUID(job_id))
-            if row is not None and row.status in (
-                JobStatus.COMPLETED,
-                JobStatus.FAILED,
-                JobStatus.TIMED_OUT,
-            ):
-                break
-            await asyncio.sleep(0.1)
-        assert row is not None, "job row was never created"
-        assert row.status == JobStatus.COMPLETED, f"job did not complete: {row.status}"
+    monkeypatch.setattr(settings, "vertex_builds_storage_enabled", False)
+    submit = await client.post("api/v2/workflows", json=_body(bg_flow), headers=_headers(created_api_key))
+    assert submit.status_code == 200, submit.text
+    job_id = submit.json()["job_id"]
 
-        # Proof 1: storage OFF => no vertex_build rows persisted for this job_id.
+    row = None
+    for _ in range(200):
         async with session_scope() as session:
-            vbs = await get_vertex_builds_by_job_id(session, job_id)
-        assert not vbs, f"vertex_builds were written despite storage OFF: {len(vbs)} rows"
-
-        # Proof 2: the durable Job.result blob carries the captured terminal outputs.
-        assert isinstance(row.result, dict), f"Job.result is not a dict: {row.result!r}"
-        assert row.result.get("outputs"), f"Job.result carried no outputs: {row.result}"
-
-        # Proof 3: GET status returns the full output, sourced from Job.result (reconstruct
-        # finds nothing with storage off, so the COMPLETED branch falls back to Job.result).
-        status = await client.get("api/v2/workflows", params={"job_id": job_id}, headers=_headers(created_api_key))
-        assert status.status_code == 200, status.text
-        body = status.json()
-        assert body["status"] == "completed"
-        assert body["outputs"], f"GET status carried no outputs with vertex_builds OFF: {body}"
-        # Proof 4: the Job.result path recovers the session_id from the terminal
-        # events (parity with the vertex-build path), so a background GET can
-        # continue the same chat thread even with vertex-build storage off.
-        assert body.get("session_id"), f"GET status lost session_id with vertex_builds OFF: {body}"
-    finally:
-        settings.vertex_builds_storage_enabled = original
+            row = await session.get(Job, UUID(job_id))
+        if row is not None and row.status in (
+            JobStatus.COMPLETED,
+            JobStatus.FAILED,
+            JobStatus.TIMED_OUT,
+        ):
+            break
+        await asyncio.sleep(0.1)
+    assert row is not None, "job row was never created"
+    assert row.status == JobStatus.COMPLETED, f"job did not complete: {row.status}"
+
+    # Proof 1: storage OFF => no vertex_build rows persisted for this job_id.
+    async with session_scope() as session:
+        vbs = await get_vertex_builds_by_job_id(session, job_id)
+    assert not vbs, f"vertex_builds were written despite storage OFF: {len(vbs)} rows"
+
+    # Proof 2: the durable Job.result blob carries the captured terminal outputs.
+    assert isinstance(row.result, dict), f"Job.result is not a dict: {row.result!r}"
+    assert row.result.get("outputs"), f"Job.result carried no outputs: {row.result}"
+
+    # Proof 3: GET status returns the full output, sourced from Job.result (reconstruct
+    # finds nothing with storage off, so the COMPLETED branch falls back to Job.result).
+    status = await client.get("api/v2/workflows", params={"job_id": job_id}, headers=_headers(created_api_key))
+    assert status.status_code == 200, status.text
+    body = status.json()
+    assert body["status"] == "completed"
+    assert body["outputs"], f"GET status carried no outputs with vertex_builds OFF: {body}"
+    # Proof 4: the Job.result path recovers session_id from the persisted submit
+    # request in job.job_metadata["request"], so a background GET can continue
+    # the same chat thread even with vertex-build storage off.
+    assert body.get("session_id"), f"GET status lost session_id with vertex_builds OFF: {body}"
 
 
 async def test_background_agui_populates_job_result_outputs(client, created_api_key, bg_flow):