From 8cba565bd2919a52b6370e37488e2d08eed0797e Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Tue, 9 Jun 2026 15:13:27 -0300 Subject: [PATCH 01/11] fix(api/v2): enforce no-code-execution gate on public workflow endpoint The v2 public endpoint only ran validate_flow_for_current_settings and skipped validate_public_flow_no_code_execution, which the v1 build_public_tmp path applies. A public flow containing a Python interpreter/REPL (or the legacy Python Code Structured tool, Smart Transform lambda) was therefore an unauthenticated server-side code-execution primitive (report H1-3754930). Mirror v1: import the validator and call it right after the public-access gate. PublicFlowValidationError subclasses CustomComponentValidationError, so the existing handler already sanitizes it to a 400 'This flow cannot be executed.' without leaking the blocked component class names. Add a non-mocking test that builds a public flow with a real PythonREPLComponent and asserts the sanitized 400 (verified RED: returns 200 without the gate). LE-1389 --- .../base/langflow/api/v2/workflow_public.py | 11 ++++- .../tests/unit/api/v2/test_workflow_public.py | 43 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/backend/base/langflow/api/v2/workflow_public.py b/src/backend/base/langflow/api/v2/workflow_public.py index 0626b4df40..52c9686eee 100644 --- a/src/backend/base/langflow/api/v2/workflow_public.py +++ b/src/backend/base/langflow/api/v2/workflow_public.py @@ -35,7 +35,11 @@ from lfx.schema.workflow import ( WORKFLOW_EXECUTION_RESPONSES, PublicWorkflowRunRequest, ) -from lfx.utils.flow_validation import CustomComponentValidationError, validate_flow_for_current_settings +from lfx.utils.flow_validation import ( + CustomComponentValidationError, + validate_flow_for_current_settings, + validate_public_flow_no_code_execution, +) from langflow.api.utils.flow_utils import ( scope_session_to_namespace, @@ -123,6 +127,11 @@ async def execute_public_workflow( flow = await session.get(Flow, real_flow_id) if flow and flow.data: validate_flow_for_current_settings(flow.data) + # Block unauthenticated execution of flows that run arbitrary code + # (Python interpreter/REPL, legacy Python Code Structured tool, + # Smart Transform lambda). Without this, any public flow containing + # such a component is an unauthenticated code-execution primitive. + validate_public_flow_no_code_execution(flow.data) flow_name = flow.name if flow else None except CustomComponentValidationError as exc: # The raw message embeds the blocked component class names; do diff --git a/src/backend/tests/unit/api/v2/test_workflow_public.py b/src/backend/tests/unit/api/v2/test_workflow_public.py index 67128fd1e3..6d1e6653b7 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_public.py +++ b/src/backend/tests/unit/api/v2/test_workflow_public.py @@ -353,6 +353,49 @@ async def test_public_endpoint_sanitizes_component_validation_error(client: Asyn assert raw_message not in response.text +@pytest.mark.benchmark +@pytest.mark.security +async def test_public_endpoint_rejects_code_execution_components( + client: AsyncClient, json_memory_chatbot_no_llm, logged_in_headers +): + """Report H1-3754930: unauthenticated public v2 runs must reject code-execution components. + + Mirrors v1 ``test_build_public_tmp_rejects_code_execution_components``. A public + flow containing a Python interpreter/REPL would otherwise let any anonymous + visitor trigger server-side code execution through ``/api/v2/workflows/public``. + """ + import json + + from tests.unit.build_utils import create_flow + + flow_dict = json.loads(json_memory_chatbot_no_llm) + flow_dict["data"]["nodes"].append( + { + "id": "PythonREPLComponent-pub1", + "type": "genericNode", + "position": {"x": 0, "y": 0}, + "data": { + "id": "PythonREPLComponent-pub1", + "type": "PythonREPLComponent", + "display_name": "Python Interpreter", + "node": {"display_name": "Python Interpreter", "template": {}}, + }, + } + ) + flow_id = await create_flow(client, json.dumps(flow_dict), logged_in_headers) + await _make_flow_public(client, flow_id, logged_in_headers) + + _send_unauthenticated(client, "test-code-exec-client") + response = await client.post( + "api/v2/workflows/public", + json={"flow_id": str(flow_id), "input_value": "Hi"}, + headers={"Content-Type": "application/json"}, + ) + + assert response.status_code == codes.BAD_REQUEST + assert response.json()["detail"] == "This flow cannot be executed." + + @pytest.mark.benchmark @pytest.mark.security async def test_public_endpoint_surfaces_value_error_as_400(client: AsyncClient, public_flow_id, monkeypatch): From 2e214d77bb256b76b2a2660f88982179c6923586 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Tue, 9 Jun 2026 16:01:04 -0300 Subject: [PATCH 02/11] fix(api/v2): reconstruct background workflow status from job-keyed vertex builds A completed background job's GET status 500'd with 'No vertex builds found for job_id'. The background build path differed from the sync path twice: 1. generate_flow_events minted a fresh run_id instead of using job_id, so vertex builds were keyed by an id the status query never uses. Thread run_id through _stream_event_frames -> generate_flow_events and pass job_id from the background buffer so graph.run_id == job_id (the sync path already does graph.set_run_id(job_id)). 2. The SSE build loop (build_vertices) only persisted builds when log_builds was set and never passed job_id. Tie log_builds to job-tracked runs (run_id present) and pass job_id=graph.run_id on the persist call. Job-tracked runs also persist streaming terminal vertices so reconstruction is complete; the live build path (run_id is None) keeps its original behavior, so the v1 build path is unchanged. Test: a real background run polled to completion, then GET status asserts a reconstructed 200 (verified RED: 500 'No vertex builds found' before the fix). Covers the non-streaming flow. v1 build path unchanged (35 build tests pass); AG-UI suite 46 pass. LE-1389 --- src/backend/base/langflow/api/build.py | 24 ++++++-- src/backend/base/langflow/api/v2/workflow.py | 10 +++- .../tests/unit/api/v2/test_workflow_agui.py | 55 +++++++++++++++++++ 3 files changed, 82 insertions(+), 7 deletions(-) diff --git a/src/backend/base/langflow/api/build.py b/src/backend/base/langflow/api/build.py index b703e72c9b..ecee6b4bed 100644 --- a/src/backend/base/langflow/api/build.py +++ b/src/backend/base/langflow/api/build.py @@ -363,6 +363,7 @@ async def generate_flow_events( current_user: CurrentActiveUser, flow_name: str | None = None, source_flow_id: uuid.UUID | None = None, + run_id: str | None = None, ) -> None: """Generate events for flow building process. @@ -370,6 +371,10 @@ async def generate_flow_events( - Building and validating the graph - Processing vertices - Handling errors and cleanup + + When ``run_id`` is provided the graph adopts it instead of minting a fresh + one, so callers (e.g. background jobs) can later look up the run's vertex + builds by that id. Defaults to a fresh uuid for the live build path. """ chat_service = get_chat_service() telemetry_service = get_telemetry_service() @@ -380,14 +385,14 @@ async def generate_flow_events( start_time = time.perf_counter() components_count = 0 graph = None - run_id = str(uuid.uuid4()) + build_run_id = run_id or str(uuid.uuid4()) try: flow_id_str = str(flow_id) # Create a fresh session for database operations async with session_scope() as fresh_session: graph = await create_graph(fresh_session, flow_id_str, flow_name) - graph.set_run_id(run_id) + graph.set_run_id(build_run_id) first_layer = sort_vertices(graph) for vertex_id in first_layer: @@ -400,13 +405,13 @@ async def generate_flow_events( vertices_to_run = list(graph.vertices_to_run.union(get_top_level_vertices(graph, graph.vertices_to_run))) await chat_service.set_cache(flow_id_str, graph) - await log_telemetry(start_time, components_count, run_id=run_id, success=True) + await log_telemetry(start_time, components_count, run_id=build_run_id, success=True) except Exception as exc: await log_telemetry( start_time, components_count, - run_id=run_id, + run_id=build_run_id, success=False, error_message=str(exc), ) @@ -523,8 +528,12 @@ async def generate_flow_events( result_data_response.message = artifacts - # Log the vertex build - if not vertex.will_stream and log_builds: + # Log the vertex build. Job-tracked runs (background workflows pass a + # ``run_id``) persist every vertex, including streaming terminal outputs, + # so GET-status reconstruction by job_id is complete. The live build + # path (``run_id is None``) keeps the original "skip streaming vertices" + # behavior unchanged. + if log_builds and (run_id is not None or not vertex.will_stream): background_tasks.add_task( log_vertex_build, flow_id=flow_id_str, @@ -533,6 +542,9 @@ async def generate_flow_events( params=params, data=result_data_response, artifacts=artifacts, + # Key the persisted build by the run id so job-tracked runs can + # reconstruct status by job_id. + job_id=graph.run_id, ) else: await chat_service.set_cache(flow_id_str, graph) diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index f156635c38..b9400124cd 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -564,6 +564,7 @@ async def _stream_event_frames( parsed: ParsedWorkflowRun, current_user: UserRead, source_flow_id: UUID | None = None, + run_id: str | None = None, ) -> AsyncIterator[tuple[bytes, str]]: """Run a flow via the v1 build-vertex loop, dispatch its events through ``adapter``. @@ -620,10 +621,14 @@ async def _stream_event_frames( files=parsed.files, stop_component_id=parsed.stop_component_id, start_component_id=parsed.start_component_id, - log_builds=False, + # Persist vertex builds (keyed by ``run_id``) only for job-tracked + # runs so a background job's status can be reconstructed later. Live + # streams pass no ``run_id`` and keep the no-persist behavior. + log_builds=run_id is not None, current_user=current_user, flow_name=flow_name, source_flow_id=source_flow_id, + run_id=run_id, ) except asyncio.CancelledError: raise @@ -922,6 +927,9 @@ async def _buffer_background_run( background_tasks=fresh_background_tasks, parsed=parsed, current_user=current_user, + # Build under the job id so the run's vertex builds are persisted + # keyed by job_id and GET-status reconstruction can find them. + run_id=job_id, ): if terminal_error_type is not None and event_type == terminal_error_type: errored = True diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index ecca272c74..7bef21a50f 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -643,6 +643,61 @@ class TestAGUIBackgroundJobStatus: "the literal 'RUN_ERROR' so the byte-substring detector would false-positive" ) + async def test_completed_background_job_status_reconstructs_from_job_id( + self, + client: AsyncClient, + created_api_key, + chatbot_flow, + ): + """A completed background job's GET status must reconstruct its outputs. + + Regression: the background build path minted its own ``run_id`` instead of + using ``job_id``, so vertex builds were persisted under a different id than + ``get_vertex_builds_by_job_id`` queries. Completed background jobs then 500 + with "No vertex builds found" on status reconstruction. The sync path + already aligns ``run_id`` to ``job_id`` (``graph.set_run_id(job_id)``); the + background path must too. + """ + import asyncio as _asyncio + from uuid import UUID as _UUID + + from langflow.services.database.models.jobs.model import Job as _Job + + headers = {"x-api-key": created_api_key.api_key} + start = await client.post( + "api/v2/workflows", + json=_agui_body(chatbot_flow, message="hi", mode="background"), + headers=headers, + ) + assert start.status_code == 200 + job_id = start.json()["job_id"] + + # Drain the SSE buffer so the build runs and persists vertex builds. + events = await client.get(f"api/v2/workflows/{job_id}/events", headers=headers) + assert events.status_code == 200 + assert "RUN_FINISHED" in events.text + + # Wait for the buffer task to finalize the job row to completed. + row = None + for _ in range(100): + async with session_scope() as session: + row = await session.get(_Job, _UUID(job_id)) + if row is not None and row.status.value in ("completed", "failed"): + break + await _asyncio.sleep(0.1) + assert row is not None, "background job row was never created" + assert row.status.value == "completed", f"background job did not complete: {row.status.value!r}" + + # GET status reconstructs from the vertex_build table keyed by job_id. + status_resp = await client.get(f"api/v2/workflows?job_id={job_id}", headers=headers) + assert status_resp.status_code == 200, ( + "completed background job status failed to reconstruct " + f"(run_id/job_id mismatch?): {status_resp.status_code} {status_resp.text}" + ) + body = status_resp.json() + assert body["status"] == "completed" + assert "outputs" in body + async def test_message_with_json_shaped_run_error_payload_does_not_fail_job( self, client: AsyncClient, From 0c8829944fa4462fd65ca44d70e4ee0c46ad458a Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 13:43:11 -0300 Subject: [PATCH 03/11] fix(api/v2): merge workflow AG-UI cancellation hardening LE-1389 --- src/backend/base/langflow/api/build.py | 3 +- src/backend/base/langflow/api/v2/workflow.py | 243 +++++++++++--- .../tests/unit/api/v2/test_workflow_agui.py | 272 +++++++++++++++- .../agui/__tests__/apply-state-delta.test.ts | 46 ++- .../API/agui/__tests__/run-agent.test.ts | 5 + .../__tests__/run-flow-bridge-e2e.test.ts | 297 +++++++++++++++++- .../agui/__tests__/run-flow-bridge.test.ts | 34 ++ .../src/controllers/API/agui/run-agent.ts | 9 +- .../controllers/API/agui/run-flow-bridge.ts | 163 +++++++++- .../src/stores/__tests__/flowStore.test.ts | 62 ++++ src/frontend/src/stores/flowStore.ts | 20 +- 11 files changed, 1088 insertions(+), 66 deletions(-) diff --git a/src/backend/base/langflow/api/build.py b/src/backend/base/langflow/api/build.py index ecee6b4bed..37b0b1ec86 100644 --- a/src/backend/base/langflow/api/build.py +++ b/src/backend/base/langflow/api/build.py @@ -364,6 +364,7 @@ async def generate_flow_events( flow_name: str | None = None, source_flow_id: uuid.UUID | None = None, run_id: str | None = None, + track_job_status: bool = True, ) -> None: """Generate events for flow building process. @@ -688,7 +689,7 @@ async def generate_flow_events( _build_run_id: uuid.UUID | None = None try: _build_run_id = uuid.UUID(graph.run_id) if graph.run_id else None - if _build_run_id is not None: + if track_job_status and _build_run_id is not None: _build_job_svc = get_job_service() await _build_job_svc.create_job( job_id=_build_run_id, diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index b9400124cd..f7a04b9e28 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -49,10 +49,10 @@ from lfx.schema.workflow import ( WorkflowStopResponse, ) from lfx.services.deps import injectable_session_scope_readonly +from lfx.utils.flow_validation import CustomComponentValidationError, validate_flow_for_current_settings from pydantic_core import ValidationError as PydanticValidationError from sqlalchemy.exc import OperationalError -from langflow.api.build import generate_flow_events from langflow.api.utils import extract_global_variables_from_headers from langflow.api.v1.schemas import FlowDataRequest, RunResponse from langflow.api.v2.adapters import ( @@ -82,6 +82,7 @@ from langflow.helpers.flow import get_flow_by_id_or_endpoint_name from langflow.processing.process import process_tweaks, run_graph_internal from langflow.services.auth.utils import get_current_user_for_workflow from langflow.services.authorization import FlowAction, ensure_flow_permission +from langflow.services.authorization.fetch import deny_to_404 from langflow.services.database.models.flow.model import FlowRead from langflow.services.database.models.jobs.model import JobType from langflow.services.database.models.user.model import UserRead @@ -94,6 +95,13 @@ EXECUTION_TIMEOUT = 300 # 5 minutes default timeout for sync execution router = APIRouter(prefix="/workflows", tags=["Workflow"]) +async def generate_flow_events(*args, **kwargs) -> None: + """Lazily call the v1 build stream to avoid import cycles during router setup.""" + from langflow.api.build import generate_flow_events as _generate_flow_events + + await _generate_flow_events(*args, **kwargs) + + def _unknown_protocol_http_exception(exc: UnknownStreamProtocolError) -> HTTPException: """Build the 422 response shared by ``stream`` and ``background`` paths. @@ -113,6 +121,66 @@ def _unknown_protocol_http_exception(exc: UnknownStreamProtocolError) -> HTTPExc ) +def _flow_not_found_privacy_exception(exc: HTTPException, flow_id: str) -> HTTPException: + return deny_to_404(exc, detail=f"Flow with id {flow_id} not found") + + +def _reject_unsupported_sync_fields(parsed: ParsedWorkflowRun) -> None: + """Reject request fields the inline sync path does not execute.""" + if parsed.mode != "sync": + return + + unsupported_fields: list[str] = [] + if parsed.data is not None: + unsupported_fields.append("data") + if parsed.files: + unsupported_fields.append("files") + if parsed.start_component_id is not None: + unsupported_fields.append("start_component_id") + if parsed.stop_component_id is not None: + unsupported_fields.append("stop_component_id") + + if unsupported_fields: + fields = ", ".join(unsupported_fields) + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "error": "Unsupported sync request fields", + "code": "SYNC_MODE_UNSUPPORTED_FIELDS", + "message": f"mode='sync' does not support request fields: {fields}. Use mode='stream' or " + "mode='background' for live-canvas overrides, files, or partial-run boundaries.", + "fields": unsupported_fields, + }, + ) + + +def _enforce_flow_data_override_owner(parsed: ParsedWorkflowRun, flow: FlowRead, current_user: UserRead) -> None: + """Only the flow owner may execute caller-supplied graph data.""" + if parsed.data is None or flow.user_id == current_user.id: + return + + raise _flow_not_found_privacy_exception( + HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Only the flow owner can override flow data during execution", + ), + parsed.flow_id, + ) + + +def _validate_flow_data_for_execution(parsed: ParsedWorkflowRun, flow: FlowRead) -> None: + """Apply the same server-side component policy gate used by v1/public runs.""" + try: + if parsed.data is not None: + validate_flow_for_current_settings(parsed.data) + elif flow.data: + validate_flow_for_current_settings(flow.data) + except CustomComponentValidationError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc + except RuntimeError as exc: + raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc)) from exc + + def _validate_output_ids(output_ids: list[str] | None, terminal_node_ids: list[str]) -> None: """Reject ``output_ids`` that aren't outputs of this flow, BEFORE it runs. @@ -221,14 +289,21 @@ async def execute_workflow( current_user.id, widen_for_shares=True, ) - await ensure_flow_permission( - current_user, - FlowAction.EXECUTE, - flow_id=flow.id, - flow_user_id=flow.user_id, - workspace_id=getattr(flow, "workspace_id", None), - folder_id=getattr(flow, "folder_id", None), - ) + try: + await ensure_flow_permission( + current_user, + FlowAction.EXECUTE, + flow_id=flow.id, + flow_user_id=flow.user_id, + workspace_id=getattr(flow, "workspace_id", None), + folder_id=getattr(flow, "folder_id", None), + ) + except HTTPException as exc: + raise _flow_not_found_privacy_exception(exc, parsed.flow_id) from exc + + _reject_unsupported_sync_fields(parsed) + _enforce_flow_data_override_owner(parsed, flow, current_user) + _validate_flow_data_for_execution(parsed, flow) if parsed.mode == "sync": return await execute_sync_workflow_with_timeout( @@ -555,6 +630,56 @@ def _single_input_value_request(parsed: ParsedWorkflowRun) -> InputValueRequest ) +_QueueItem = tuple[str | None, bytes | None, float] + + +class _WorkflowEventQueue: + """Bounded EventManager handoff that fails explicitly instead of dropping events.""" + + def __init__(self, maxsize: int) -> None: + self._queue: asyncio.Queue[_QueueItem] = asyncio.Queue(maxsize=maxsize) + self._overflowed = False + self._loop = asyncio.get_running_loop() + self._overflow_task: asyncio.Task[None] | None = None + + @property + def maxsize(self) -> int: + return self._queue.maxsize + + async def get(self) -> _QueueItem: + return await self._queue.get() + + async def put(self, item: _QueueItem) -> None: + if self._overflowed: + return + await self._queue.put(item) + + def put_nowait(self, item: _QueueItem) -> None: + if self._overflowed: + return + try: + self._queue.put_nowait(item) + except asyncio.QueueFull: + self._overflowed = True + self._overflow_task = self._loop.create_task(self._emit_overflow_error()) + + async def _emit_overflow_error(self) -> None: + payload = { + "event": "error", + "data": { + "error": "Workflow event stream exceeded buffering capacity; client is consuming events too slowly." + }, + } + await self._queue.put((f"error-{uuid4()}", json.dumps(payload).encode("utf-8"), time.time())) + await self._queue.put((None, None, time.time())) + + async def aclose(self) -> None: + if self._overflow_task is not None and not self._overflow_task.done(): + self._overflow_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._overflow_task + + async def _stream_event_frames( *, adapter: StreamAdapter, @@ -565,6 +690,7 @@ async def _stream_event_frames( current_user: UserRead, source_flow_id: UUID | None = None, run_id: str | None = None, + track_job_status: bool = True, ) -> AsyncIterator[tuple[bytes, str]]: """Run a flow via the v1 build-vertex loop, dispatch its events through ``adapter``. @@ -583,10 +709,10 @@ async def _stream_event_frames( playground's chat-view. A follow-up retires this once chat-view consumes the AG-UI ``TEXT_MESSAGE_*`` lifecycle directly. """ - # Bounded so a slow consumer can apply backpressure on the build loop - # instead of growing without bound. The build loop awaits ``queue.put`` - # which yields control back to the consumer between frames. - queue: asyncio.Queue = asyncio.Queue(maxsize=_EVENT_QUEUE_MAX_SIZE) + # EventManager uses put_nowait(), so a plain bounded asyncio.Queue would + # silently drop frames via QueueFull. This adapter keeps memory bounded and + # converts overflow into an explicit stream error + sentinel. + queue = _WorkflowEventQueue(maxsize=_EVENT_QUEUE_MAX_SIZE) event_manager = create_default_event_manager(queue) input_request = _single_input_value_request(parsed) flow_data = FlowDataRequest(**parsed.data) if parsed.data else None @@ -629,6 +755,7 @@ async def _stream_event_frames( flow_name=flow_name, source_flow_id=source_flow_id, run_id=run_id, + track_job_status=track_job_status, ) except asyncio.CancelledError: raise @@ -650,7 +777,7 @@ async def _stream_event_frames( # side-channel ``CustomEvent``; emitted only when the wire protocol is # AG-UI. A follow-up retires this once chat-view consumes AG-UI primitives. emit_side_channel = adapter.name == "agui" - side_channel_events = frozenset({"add_message", "token", "remove_message", "error"}) + side_channel_events = frozenset({"add_message", "token", "remove_message", "error", "end"}) seq = 0 run_task = asyncio.create_task(drive()) @@ -698,6 +825,7 @@ async def _stream_event_frames( run_task.cancel() with contextlib.suppress(asyncio.CancelledError): await run_task + await queue.aclose() def _execute_streaming_workflow( @@ -930,6 +1058,7 @@ async def _buffer_background_run( # Build under the job id so the run's vertex builds are persisted # keyed by job_id and GET-status reconstruction can find them. run_id=job_id, + track_job_status=False, ): if terminal_error_type is not None and event_type == terminal_error_type: errored = True @@ -948,20 +1077,6 @@ async def _buffer_background_run( # passing the raw string would silently miss every row. job_uuid = UUID(job_id) if isinstance(job_id, str) else job_id await _finalize_job_status(job_uuid, JobStatus.FAILED if errored else JobStatus.COMPLETED) - # Fire memory-base auto-capture hook on successful runs only. Matches - # the sync mode wiring above and the v1 build-pipeline wiring in - # ``api/build.py``. ``fire_and_forget_task`` because we are already a - # background coroutine and the hook must not block job finalization. - if not errored: - try: - await get_task_service().fire_and_forget_task( - get_memory_base_service().on_flow_output, - flow_id=flow.id, - session_id=parsed.session_id or str(flow.id), - job_id=job_uuid, - ) - except (RuntimeError, ValueError, OSError): - await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True) async def execute_workflow_background( @@ -997,6 +1112,7 @@ async def execute_workflow_background( try: queue_service = get_queue_service() queue_service.create_queue(job_id_str) + await queue_service.register_job_owner(job_id_str, current_user.id) queue_service.start_job( job_id_str, _buffer_background_run( @@ -1215,7 +1331,6 @@ async def stop_workflow( """ job_id = request.job_id job_service = get_job_service() - task_service = get_task_service() try: # 1. Fetch Job @@ -1257,27 +1372,29 @@ async def stop_workflow( return WorkflowStopResponse(job_id=str(job_id), message=f"Job {job_id} is already cancelled.") try: - revoked = await task_service.revoke_task(job_id) + try: + await get_queue_service().cancel_job(str(job_id)) + except asyncio.CancelledError as exc: + message_code = exc.args[0] if exc.args else "UNKNOWN" + if message_code != "LANGFLOW_USER_CANCELLED": + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "error": "Task cancellation error", + "code": message_code, + "message": f"Job {job_id} was cancelled unexpectedly by the system", + "job_id": str(job_id), + }, + ) from exc await job_service.update_job_status(job_id, JobStatus.CANCELLED) # Release the in-memory buffer and wake any re-attach waiters so they # see a clean stream-end instead of hanging on ``_cond.wait()``. await _clear_background_run(str(job_id)) - message = f"Job {job_id} cancelled successfully." if revoked else f"Job {job_id} is already cancelled." + message = f"Job {job_id} cancelled successfully." return WorkflowStopResponse(job_id=str(job_id), message=message) - except asyncio.CancelledError as exc: - # Handle system-initiated cancellations that were re-raised - # The job status has already been updated to FAILED in jobs/service.py - message_code = exc.args[0] if exc.args else "UNKNOWN" - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={ - "error": "Task cancellation error", - "code": message_code, - "message": f"Job {job_id} was cancelled unexpectedly by the system", - "job_id": str(job_id), - }, - ) from exc + except HTTPException: + raise except Exception as exc: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -1308,7 +1425,7 @@ async def reattach_workflow_events( leaking the existence of other users' runs. """ bg_run = _BACKGROUND_RUNS.get(job_id) - if bg_run is None or bg_run.user_id != str(current_user.id): + if bg_run is not None and bg_run.user_id != str(current_user.id): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail={ @@ -1319,6 +1436,42 @@ async def reattach_workflow_events( }, ) + if bg_run is None: + try: + job_uuid = UUID(job_id) + except ValueError: + job_uuid = None + + job = None + if job_uuid is not None: + with contextlib.suppress(Exception): + job = await get_job_service().get_job_by_job_id(job_uuid, user_id=current_user.id) + + if job is None or job.type != JobType.WORKFLOW: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={ + "error": "Background run not found", + "code": "JOB_NOT_FOUND", + "message": f"No buffered events for job {job_id}.", + "job_id": job_id, + }, + ) + + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": "Background event buffer unavailable", + "code": "BACKGROUND_EVENTS_UNAVAILABLE", + "message": ( + f"Buffered events for job {job_id} are not available on this worker. " + "Use the workflow status endpoint for this job, or route event re-attach " + "requests to the worker that accepted the background run." + ), + "job_id": job_id, + }, + ) + last_event_id = http_request.headers.get("Last-Event-ID") start_index = 0 if last_event_id: diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index 7bef21a50f..f9f6d60c1b 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -11,10 +11,16 @@ Execution mode (``mode`` field): - ``stream``: SSE in the chosen ``stream_protocol``. """ +import asyncio +import contextlib import json +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock from uuid import uuid4 import pytest +from fastapi import HTTPException from httpx import AsyncClient from langflow.services.database.models.flow.model import Flow from lfx.services.deps import session_scope @@ -222,9 +228,143 @@ class TestAGUIModeDispatch: assert result["status"] == "completed" +class TestV2WorkflowAdmission: + """Route-level admission checks before workflow execution dispatch.""" + + async def test_non_owner_data_override_is_hidden_as_404(self, monkeypatch: pytest.MonkeyPatch): + """Execute-only sharees must not inject alternate graph data into shared flows.""" + from langflow.api.v2 import workflow as workflow_module + from lfx.schema.workflow import WorkflowRunRequest + + flow_id = uuid4() + owner_id = uuid4() + caller_id = uuid4() + flow = SimpleNamespace( + id=flow_id, + user_id=owner_id, + workspace_id=None, + folder_id=None, + data={"nodes": [], "edges": []}, + name="shared", + ) + current_user = SimpleNamespace(id=caller_id) + + monkeypatch.setattr(workflow_module, "get_flow_by_id_or_endpoint_name", AsyncMock(return_value=flow)) + monkeypatch.setattr(workflow_module, "ensure_flow_permission", AsyncMock(return_value=None)) + + with pytest.raises(HTTPException) as exc_info: + await workflow_module.execute_workflow( + WorkflowRunRequest( + flow_id=str(flow_id), + input_value="hi", + mode="stream", + data={"nodes": [], "edges": []}, + ), + background_tasks=SimpleNamespace(), + http_request=SimpleNamespace(), + current_user=current_user, + ) + + assert exc_info.value.status_code == 404 + assert exc_info.value.detail["code"] == "FLOW_NOT_FOUND" + + async def test_execute_permission_denial_is_hidden_as_404(self, monkeypatch: pytest.MonkeyPatch): + """A denied share-aware fetch must not leak flow existence as a raw 403.""" + from langflow.api.v2 import workflow as workflow_module + from lfx.schema.workflow import WorkflowRunRequest + + flow_id = uuid4() + caller_id = uuid4() + flow = SimpleNamespace( + id=flow_id, + user_id=uuid4(), + workspace_id=None, + folder_id=None, + data={"nodes": [], "edges": []}, + name="private", + ) + + async def _deny(*_args, **_kwargs): + raise HTTPException(status_code=403, detail="denied") + + monkeypatch.setattr(workflow_module, "get_flow_by_id_or_endpoint_name", AsyncMock(return_value=flow)) + monkeypatch.setattr(workflow_module, "ensure_flow_permission", _deny) + + with pytest.raises(HTTPException) as exc_info: + await workflow_module.execute_workflow( + WorkflowRunRequest(flow_id=str(flow_id), input_value="hi", mode="stream"), + background_tasks=SimpleNamespace(), + http_request=SimpleNamespace(), + current_user=SimpleNamespace(id=caller_id), + ) + + assert exc_info.value.status_code == 404 + assert exc_info.value.detail["code"] == "FLOW_NOT_FOUND" + + class TestAGUIStreaming: """mode=stream returns an AG-UI server-sent event stream.""" + async def test_stream_event_queue_ignores_late_sentinel_after_overflow(self): + """A normal completion sentinel must not race ahead of the overflow error.""" + from langflow.api.v2 import workflow as workflow_module + + queue = workflow_module._WorkflowEventQueue(maxsize=1) + try: + payload = json.dumps({"event": "token", "data": {"chunk": "0"}}).encode() + queue.put_nowait(("token-0", payload, time.time())) + queue.put_nowait(("token-1", payload, time.time())) + + await asyncio.wait_for(queue.put((None, None, time.time())), timeout=0.1) + finally: + await queue.aclose() + + async def test_stream_event_handoff_overflow_emits_error(self, monkeypatch: pytest.MonkeyPatch): + """EventManager put_nowait must not silently drop frames when the stream buffer fills.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.adapters import StreamEvent + from langflow.api.v2.converters import ParsedWorkflowRun + + seen_maxsize: list[int] = [] + + async def fake_generate_flow_events(**kwargs): + event_queue = kwargs["event_manager"].queue + seen_maxsize.append(event_queue.maxsize) + for index in range(event_queue.maxsize + 1): + payload = json.dumps({"event": "token", "data": {"chunk": str(index)}}).encode() + event_queue.put_nowait((f"token-{index}", payload, time.time())) + await event_queue.put((None, None, time.time())) + + class FakeAdapter: + name = "langflow" + + def initial_events(self): + return [] + + def final_events(self): + return [] + + def translate(self, event_type, _event_data): + return [StreamEvent(type=event_type, data_json="{}")] + + monkeypatch.setattr(workflow_module, "_EVENT_QUEUE_MAX_SIZE", 2) + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + event_types = [ + event_type + async for _frame, event_type in workflow_module._stream_event_frames( + adapter=FakeAdapter(), + flow_id=uuid4(), + flow_name="flow", + background_tasks=SimpleNamespace(add_task=lambda *_args, **_kwargs: None), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="", mode="stream"), + current_user=SimpleNamespace(id=uuid4()), + ) + ] + + assert seen_maxsize == [2] + assert event_types == ["token", "token", "error"] + async def test_stream_emits_run_lifecycle_events( self, client: AsyncClient, @@ -498,6 +638,32 @@ class TestAGUICancellation: class TestAGUIBackgroundReattach: """A background run buffers its AG-UI events so a client can re-attach.""" + async def test_reattach_owned_job_without_local_buffer_returns_409(self, monkeypatch: pytest.MonkeyPatch): + """If the job exists but this worker has no replay buffer, return an explicit conflict.""" + from langflow.api.v2 import workflow as workflow_module + + job_id = uuid4() + expected_user_id = uuid4() + workflow_module._BACKGROUND_RUNS.pop(str(job_id), None) + + class FakeJobService: + async def get_job_by_job_id(self, seen_job_id, user_id=None): + assert seen_job_id == job_id + assert user_id == expected_user_id + return SimpleNamespace(type=workflow_module.JobType.WORKFLOW) + + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + + with pytest.raises(HTTPException) as exc_info: + await workflow_module.reattach_workflow_events( + str(job_id), + http_request=SimpleNamespace(headers={}), + current_user=SimpleNamespace(id=expected_user_id), + ) + + assert exc_info.value.status_code == 409 + assert exc_info.value.detail["code"] == "BACKGROUND_EVENTS_UNAVAILABLE" + async def test_background_run_events_can_be_reattached( self, client: AsyncClient, @@ -594,6 +760,38 @@ class TestAGUIBackgroundJobStatus: payload that happened to contain the literal byte sequence ``"RUN_ERROR"``. """ + async def test_background_buffer_binds_build_rows_to_returned_job_id(self, monkeypatch: pytest.MonkeyPatch): + """Status reconstruction needs vertex_build rows logged under the public background job id.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + job_id = uuid4() + captured: dict = {} + + class FakeAdapter: + terminal_error_type = "RUN_ERROR" + + async def fake_stream_event_frames(**kwargs): + captured.update(kwargs) + yield b"data: {}\n\n", "RUN_FINISHED" + + monkeypatch.setattr(workflow_module, "get_stream_adapter", lambda *_args, **_kwargs: FakeAdapter()) + monkeypatch.setattr(workflow_module, "_stream_event_frames", fake_stream_event_frames) + monkeypatch.setattr(workflow_module, "_finalize_job_status", AsyncMock()) + + bg_run = workflow_module._BackgroundRun(user_id=str(uuid4())) + await workflow_module._buffer_background_run( + bg_run=bg_run, + flow=SimpleNamespace(id=uuid4(), name="flow"), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="hi", mode="background"), + job_id=str(job_id), + current_user=SimpleNamespace(id=uuid4()), + stream_protocol="agui", + ) + + assert captured["run_id"] == str(job_id) + assert captured["track_job_status"] is False + async def test_message_text_containing_run_error_literal_does_not_fail_job( self, client: AsyncClient, @@ -988,11 +1186,10 @@ class TestMemoryBaseHookBackgroundMode: """The memory-base ``on_flow_output`` hook must fire after a background run. Sync mode wires the hook directly inside ``execute_workflow_sync`` and the - v1 build pipeline wires it after ``end_all_traces`` in ``api/build.py``. - The v2 background mode buffers frames in ``_buffer_background_run`` and - must dispatch the same hook in its ``finally`` block on successful - completion. Without it, MemoryBase auto-capture silently misses every - background run. + v1 build pipeline wires it after ``end_all_traces`` in ``api/build.py``. The + v2 background mode routes through that same build pipeline with the returned + background job id, so the hook must fire exactly once for successful + background runs. """ async def test_background_run_fires_memory_base_hook_on_success( @@ -1006,7 +1203,7 @@ class TestMemoryBaseHookBackgroundMode: ``flow_id``, ``session_id``, and ``job_id`` must reach the hook. """ - from langflow.api.v2 import workflow as _workflow_module + from langflow.api import build as _build_module captured: list[dict] = [] @@ -1015,7 +1212,7 @@ class TestMemoryBaseHookBackgroundMode: captured.append(kwargs) monkeypatch.setattr( - _workflow_module, + _build_module, "get_memory_base_service", lambda: _RecordingMemoryBaseService(), ) @@ -1643,3 +1840,64 @@ class TestStopWorkflowEndToEnd: row = await session.get(Job, _UUID(job_id)) assert row is not None assert row.status == JobStatus.CANCELLED + + async def test_stop_cancels_queue_owned_background_task_when_task_service_noops( + self, + client: AsyncClient, + created_api_key, + chatbot_flow, + monkeypatch, + ): + """The v2 background runner is owned by the queue service, not TaskService.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.services.deps import get_queue_service + from langflow.services.job_queue.service import JobQueueNotFoundError + + started = asyncio.Event() + cancelled = asyncio.Event() + job_id: str | None = None + + async def _never_finishes(**_kwargs): + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + class NoopTaskService: + async def revoke_task(self, _task_id): + return True + + monkeypatch.setattr(workflow_module, "_buffer_background_run", _never_finishes) + monkeypatch.setattr(workflow_module, "get_task_service", lambda: NoopTaskService()) + + headers = {"x-api-key": created_api_key.api_key} + try: + start = await client.post( + "api/v2/workflows", + json=_agui_body(chatbot_flow, mode="background"), + headers=headers, + ) + assert start.status_code == 200 + job_id = start.json()["job_id"] + await asyncio.wait_for(started.wait(), timeout=2) + + stop = await client.post( + "api/v2/workflows/stop", + json={"job_id": job_id}, + headers=headers, + ) + + assert stop.status_code == 200 + assert cancelled.is_set() + try: + _, _, task, _ = get_queue_service().get_queue_data(job_id) + except JobQueueNotFoundError: + task = None + assert task is None or task.done() + finally: + if job_id is not None: + with contextlib.suppress(BaseException): + await get_queue_service().cleanup_job(job_id) + await workflow_module._clear_background_run(job_id) diff --git a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts index dc7a471903..b6a7bd1427 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts @@ -218,7 +218,8 @@ describe("applyStateDelta", () => { * The v1 build path drove edge animation through onBuildStart/onBuildEnd. * Without an explicit toggle here the AG-UI path used to leave edges * static until ``finish()`` cleared them. Pin the contract: running flips - * edges on, success/error flips them off (per-node, no global resets). + * edges on, success/error flips them off, and the production caller keeps + * the full active node set animated across parallel branches. */ it("flips edges on when a node enters running status", () => { const touched = new Set(); @@ -281,5 +282,48 @@ describe("applyStateDelta", () => { { ids: ["node-b"], running: false }, ]); }); + + it("keeps all currently running node edges animated together", () => { + const touched = new Set(); + const running = new Set(); + const calls: string[][] = []; + const original = useFlowStore.getState().clearAndSetEdgesRunning; + useFlowStore.setState({ + clearAndSetEdgesRunning: (ids?: string[]) => { + calls.push(ids ?? []); + }, + }); + + applyStateDelta( + [ + { + op: "add", + path: "/nodes/node-a", + value: { status: "running", output: null }, + }, + { + op: "add", + path: "/nodes/node-b", + value: { status: "running", output: null }, + }, + { + op: "add", + path: "/nodes/node-a", + value: { status: "success", output: { results: {} } }, + }, + ], + "run-1", + touched, + running, + ); + + useFlowStore.setState({ clearAndSetEdgesRunning: original }); + + expect(calls).toEqual([ + ["node-a"], + ["node-a", "node-b"], + ["node-b"], + ]); + }); }); }); diff --git a/src/frontend/src/controllers/API/agui/__tests__/run-agent.test.ts b/src/frontend/src/controllers/API/agui/__tests__/run-agent.test.ts index 0540aefbb7..9ea94dd08e 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/run-agent.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/run-agent.test.ts @@ -5,6 +5,10 @@ import { WORKFLOWS_PUBLIC_ENDPOINT, } from "../run-agent"; +jest.mock("@/customization/utils/get-fetch-credentials", () => ({ + getFetchCredentials: jest.fn(() => "include"), +})); + describe("buildWorkflowRunRequest", () => { it("returns a native WorkflowRunRequest with input_value mapped from message", () => { const body = buildWorkflowRunRequest({ @@ -268,6 +272,7 @@ describe("createWorkflowAgent wire body", () => { expect(captured).toHaveLength(1); expect(captured[0].url).toBe(WORKFLOWS_ENDPOINT); expect(captured[0].init.method).toBe("POST"); + expect(captured[0].init.credentials).toBe("include"); const wireBody = JSON.parse(captured[0].init.body as string); expect(wireBody).toEqual({ flow_id: "67ccd2be-17f0-8190-81ff-3bb2cf6508e6", diff --git a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts index 529cc8616b..211f68ff3c 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts @@ -9,10 +9,12 @@ */ import { BuildStatus } from "@/constants/enums"; +import { queryClient } from "@/contexts"; import { runFlowAGUI } from "@/controllers/API/agui/run-flow-bridge"; import useAlertStore from "@/stores/alertStore"; import useFlowStore from "@/stores/flowStore"; import { useMessagesStore } from "@/stores/messagesStore"; +import type { Message } from "@/types/messages"; beforeAll(() => { // jsdom lacks the codecs the AG-UI SSE decoder uses; pull them from @@ -34,6 +36,7 @@ beforeAll(() => { }); beforeEach(() => { + queryClient.clear(); // Reset only the slices the bridge writes to. Touching the full state // would clobber action implementations (Zustand stores actions in the // same state object). @@ -42,6 +45,7 @@ beforeEach(() => { flowPool: {}, isBuilding: true, buildInfo: null, + buildDuration: null, }); useAlertStore.setState({ errorData: { title: "", list: [] }, @@ -232,6 +236,240 @@ describe("runFlowAGUI end-to-end", () => { } }); + it("marks still-running nodes as ERROR when a run-level RUN_ERROR arrives", async () => { + const events = [ + { + type: "RUN_STARTED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + { + type: "STATE_DELTA", + delta: [ + { + op: "add", + path: "/nodes/node-a", + value: { status: "running", output: null }, + }, + ], + timestamp: 0, + }, + { + type: "RUN_ERROR", + message: "graph blew up", + timestamp: 0, + }, + ]; + const fetchSpy = jest + .spyOn(global as { fetch: typeof fetch }, "fetch") + .mockImplementation((async () => + sseResponse(events)) as unknown as typeof fetch); + + try { + await runFlowAGUI({ flowId: "flow-1", message: "boom" }); + + const flow = useFlowStore.getState(); + expect(flow.flowBuildStatus["node-a"]?.status).toBe(BuildStatus.ERROR); + expect(flow.buildInfo).toEqual({ + error: ["graph blew up"], + success: false, + }); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("does not write success buildInfo for silent runs", async () => { + const fetchSpy = jest + .spyOn(global as { fetch: typeof fetch }, "fetch") + .mockImplementation((async () => + sseResponse([ + { + type: "RUN_STARTED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + { + type: "RUN_FINISHED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + ])) as unknown as typeof fetch); + + try { + await runFlowAGUI({ flowId: "flow-1", message: "hi", silent: true }); + + const flow = useFlowStore.getState(); + expect(flow.isBuilding).toBe(false); + expect(flow.buildInfo).toBeNull(); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("preserves build_duration from the legacy end side-channel", async () => { + const queryKey = [ + "useGetMessagesQuery", + { id: "flow-1", session_id: "thread-1" }, + ]; + const botMessage: Message = { + id: "msg-duration", + flow_id: "flow-1", + session_id: "thread-1", + sender: "Machine", + sender_name: "agent", + text: "hello", + timestamp: "2026-06-10T12:00:00.000Z", + files: [], + edit: false, + background_color: "", + text_color: "", + properties: {}, + }; + queryClient.setQueryData(queryKey, [botMessage]); + useMessagesStore.setState({ messages: [botMessage] }); + useFlowStore.setState({ + buildingFlowId: "flow-1", + buildingSessionId: "thread-1", + }); + + const fetchSpy = jest + .spyOn(global as { fetch: typeof fetch }, "fetch") + .mockImplementation((async () => + sseResponse([ + { + type: "CUSTOM", + name: "langflow.event", + value: { event_type: "end", data: { build_duration: 1.5 } }, + timestamp: 0, + rawEvent: {}, + }, + { + type: "RUN_FINISHED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + ])) as unknown as typeof fetch); + + try { + await runFlowAGUI({ flowId: "flow-1", message: "hi" }); + + expect(useFlowStore.getState().buildDuration).toBe(1500); + expect( + queryClient.getQueryData(queryKey)?.[0].properties + ?.build_duration, + ).toBe(1500); + expect( + useMessagesStore.getState().messages[0].properties?.build_duration, + ).toBe(1500); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("does not overwrite an existing message build_duration", async () => { + const queryKey = [ + "useGetMessagesQuery", + { id: "flow-1", session_id: "thread-1" }, + ]; + const botMessage: Message = { + id: "msg-duration", + flow_id: "flow-1", + session_id: "thread-1", + sender: "Machine", + sender_name: "agent", + text: "hello", + timestamp: "2026-06-10T12:00:00.000Z", + files: [], + edit: false, + background_color: "", + text_color: "", + properties: { build_duration: 500 }, + }; + queryClient.setQueryData(queryKey, [botMessage]); + useMessagesStore.setState({ messages: [botMessage] }); + useFlowStore.setState({ + buildingFlowId: "flow-1", + buildingSessionId: "thread-1", + }); + + const fetchSpy = jest + .spyOn(global as { fetch: typeof fetch }, "fetch") + .mockImplementation((async () => + sseResponse([ + { + type: "CUSTOM", + name: "langflow.event", + value: { event_type: "end", data: { build_duration: 1.5 } }, + timestamp: 0, + rawEvent: {}, + }, + { + type: "RUN_FINISHED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + ])) as unknown as typeof fetch); + + try { + await runFlowAGUI({ flowId: "flow-1", message: "hi" }); + + expect(useFlowStore.getState().buildDuration).toBe(1500); + expect( + queryClient.getQueryData(queryKey)?.[0].properties + ?.build_duration, + ).toBe(500); + expect( + useMessagesStore.getState().messages[0].properties?.build_duration, + ).toBe(500); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("appends component log events to the flow pool", async () => { + const fetchSpy = jest + .spyOn(global as { fetch: typeof fetch }, "fetch") + .mockImplementation((async () => + sseResponse([ + { + type: "CUSTOM", + name: "langflow.log", + value: { + component_id: "node-a", + output: "debug", + name: "Debug", + message: "hello", + type: "info", + }, + timestamp: 0, + rawEvent: {}, + }, + { + type: "RUN_FINISHED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + ])) as unknown as typeof fetch); + + try { + await runFlowAGUI({ flowId: "flow-1", message: "hi" }); + + const entry = useFlowStore.getState().flowPool["node-a"]?.[0]; + expect(entry?.data.logs?.debug).toEqual([ + { name: "Debug", message: "hello", type: "info" }, + ]); + } finally { + fetchSpy.mockRestore(); + } + }); + it("records a failure when the stream closes without RUN_FINISHED or RUN_ERROR", async () => { /** * Server-side crash, truncated proxy response, or a buggy stream can @@ -275,10 +513,50 @@ describe("runFlowAGUI end-to-end", () => { * stream would keep running on the wire. */ const signals: AbortSignal[] = []; + const encoder = new ( + globalThis as { TextEncoder: typeof TextEncoder } + ).TextEncoder(); + const payload = encoder.encode( + [ + { + type: "RUN_STARTED", + threadId: "thread-1", + runId: "run-1", + timestamp: 0, + }, + { + type: "STATE_DELTA", + delta: [ + { + op: "add", + path: "/nodes/node-a", + value: { status: "running", output: null }, + }, + { + op: "add", + path: "/nodes/node-b", + value: { status: "running", output: null }, + }, + { + op: "replace", + path: "/nodes/node-b", + value: { + status: "success", + output: { results: { text: "done" } }, + }, + }, + ], + timestamp: 0, + }, + ] + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""), + ); const fetchSpy = jest .spyOn(global as { fetch: typeof fetch }, "fetch") .mockImplementation(((_input: RequestInfo | URL, init?: RequestInit) => { if (init?.signal) signals.push(init.signal); + let delivered = false; return Promise.resolve({ ok: true, status: 200, @@ -292,7 +570,13 @@ describe("runFlowAGUI end-to-end", () => { body: { getReader() { return { - read: () => new Promise(() => {}), + read: () => { + if (!delivered) { + delivered = true; + return Promise.resolve({ done: false, value: payload }); + } + return new Promise(() => {}); + }, cancel: async () => {}, }; }, @@ -302,6 +586,13 @@ describe("runFlowAGUI end-to-end", () => { }) as unknown as typeof fetch); try { + useFlowStore.setState({ + flowBuildStatus: { + "node-a": { status: BuildStatus.TO_BUILD }, + "node-b": { status: BuildStatus.TO_BUILD }, + }, + }); + const controller = new AbortController(); const runPromise = runFlowAGUI({ flowId: "flow-1", @@ -321,6 +612,10 @@ describe("runFlowAGUI end-to-end", () => { expect(signals[0].aborted).toBe(true); const flow = useFlowStore.getState(); expect(flow.isBuilding).toBe(false); + expect(flow.flowBuildStatus["node-a"]?.status).toBe( + BuildStatus.TO_BUILD, + ); + expect(flow.flowBuildStatus["node-b"]?.status).toBe(BuildStatus.BUILT); // ``buildInfo`` is recorded as a failure (so trackFlowBuild sees the // cancellation as an error) but without an ``error`` string: the // caller's stopBuilding already shows the user-facing "Build stopped" diff --git a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge.test.ts b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge.test.ts index f61cd9f432..6828f6320c 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge.test.ts @@ -20,6 +20,8 @@ function makeRecordingContext() { setRunId: (runId) => calls.push(`setRunId:${runId}`), applyDelta: (ops) => calls.push(`applyDelta:${ops.length}`), handleCustomEvent: (eventType) => calls.push(`custom:${eventType}`), + handleEndEvent: () => calls.push("end"), + handleLogEvent: () => calls.push("log"), onFinished: () => calls.push("finished"), onError: (message) => calls.push(`error:${message}`), }; @@ -123,6 +125,38 @@ describe("handleAGUIEvent non-terminal contract", () => { expect(calls).toEqual(["custom:add_message"]); }); + it("routes CUSTOM(langflow.event:end) to the end-event handler", () => { + const { ctx, calls } = makeRecordingContext(); + + const terminal = handleAGUIEvent( + { + type: EventType.CUSTOM, + name: "langflow.event", + value: { event_type: "end", data: { build_duration: 1.25 } }, + } as unknown as BaseEvent, + ctx, + ); + + expect(terminal).toBe(false); + expect(calls).toEqual(["end"]); + }); + + it("routes CUSTOM(langflow.log) to the log handler", () => { + const { ctx, calls } = makeRecordingContext(); + + const terminal = handleAGUIEvent( + { + type: EventType.CUSTOM, + name: "langflow.log", + value: { component_id: "node-a", output: "out" }, + } as unknown as BaseEvent, + ctx, + ); + + expect(terminal).toBe(false); + expect(calls).toEqual(["log"]); + }); + it("returns false for CUSTOM with a foreign name and skips the handler", () => { const { ctx, calls } = makeRecordingContext(); diff --git a/src/frontend/src/controllers/API/agui/run-agent.ts b/src/frontend/src/controllers/API/agui/run-agent.ts index 0f6e573053..d9ed04e903 100644 --- a/src/frontend/src/controllers/API/agui/run-agent.ts +++ b/src/frontend/src/controllers/API/agui/run-agent.ts @@ -16,6 +16,8 @@ import { type HttpAgentConfig, type RunAgentInput, } from "@ag-ui/client"; +import { BASE_URL_API_V2 } from "@/constants/constants"; +import { getFetchCredentials } from "@/customization/utils/get-fetch-credentials"; /** Execution mode accepted by the v2 workflows endpoint. */ export type WorkflowMode = "sync" | "stream" | "background"; @@ -43,6 +45,8 @@ export interface WorkflowRunOptions { flowData?: { nodes: unknown[]; edges: unknown[] }; /** Runtime file references the graph build needs (e.g. uploaded file paths). */ files?: string[]; + /** Suppress success build UI for chat/playground runs that should not show a build panel. */ + silent?: boolean; /** * When true, route the request to the public-flow endpoint * (``/api/v2/workflows/public``) instead of the authenticated one. @@ -60,10 +64,10 @@ export interface WorkflowRunOptions { } /** The v2 workflows endpoint path. */ -export const WORKFLOWS_ENDPOINT = "/api/v2/workflows"; +export const WORKFLOWS_ENDPOINT = `${BASE_URL_API_V2}workflows`; /** The v2 public workflows endpoint path (shareable playground). */ -export const WORKFLOWS_PUBLIC_ENDPOINT = "/api/v2/workflows/public"; +export const WORKFLOWS_PUBLIC_ENDPOINT = `${BASE_URL_API_V2}workflows/public`; /** * Wire shape for `POST /api/v2/workflows`. Mirrors @@ -185,6 +189,7 @@ export function createWorkflowAgent( }, body: JSON.stringify(agent.workflowBody), signal: agent.abortController.signal, + credentials: getFetchCredentials(), }; }; return agent; diff --git a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts index 74ed69b1f9..5dc576c6c8 100644 --- a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts +++ b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts @@ -10,12 +10,19 @@ import { type BaseEvent, EventType } from "@ag-ui/client"; import { handleMessageEvent } from "@/components/core/playgroundComponent/chat-view/utils/message-event-handler"; +import { + findLastBotMessage, + updateMessageProperties, +} from "@/components/core/playgroundComponent/chat-view/utils/message-utils"; import { BuildStatus } from "@/constants/enums"; +import { persistMessageProperties } from "@/controllers/API/helpers/persist-message-properties"; import useAlertStore from "@/stores/alertStore"; import useFlowStore from "@/stores/flowStore"; +import { useMessagesStore } from "@/stores/messagesStore"; import type { ChatInputType, ChatOutputType, + LogsLogType, VertexBuildTypeAPI, VertexDataTypeAPI, } from "@/types/api"; @@ -44,6 +51,11 @@ interface AGUINodeState { output: VertexDataTypeAPI | null; } +type FlowBuildStatusEntry = { + status: BuildStatus; + timestamp?: string; +}; + /** * Hooks the bridge exposes to {@link handleAGUIEvent}. Injecting them keeps * the event-dispatch logic testable without importing the real flow / alert @@ -53,6 +65,8 @@ export interface BridgeContext { setRunId: (runId: string) => void; applyDelta: (ops: JsonPatchOp[]) => void; handleCustomEvent: (eventType: string, data: unknown) => void; + handleEndEvent: (data: unknown) => void; + handleLogEvent: (data: unknown) => void; onFinished: () => void; onError: (message: string) => void; } @@ -82,7 +96,13 @@ export function handleAGUIEvent(event: BaseEvent, ctx: BridgeContext): boolean { value?: { event_type?: string; data?: unknown }; }; if (custom.name === "langflow.event" && custom.value?.event_type) { - ctx.handleCustomEvent(custom.value.event_type, custom.value.data); + if (custom.value.event_type === "end") { + ctx.handleEndEvent(custom.value.data); + } else { + ctx.handleCustomEvent(custom.value.event_type, custom.value.data); + } + } else if (custom.name === "langflow.log") { + ctx.handleLogEvent(custom.value); } return false; } @@ -112,6 +132,8 @@ export function applyStateDelta( ops: JsonPatchOp[], runId: string, nodeIds: Set, + runningNodeIds?: Set, + originalBuildStatuses?: Map, ): void { const flowStore = useFlowStore.getState(); for (const op of ops) { @@ -124,6 +146,13 @@ export function applyStateDelta( const buildStatus = AGUI_STATUS_TO_BUILD_STATUS[value.status] ?? BuildStatus.BUILDING; + if (originalBuildStatuses && !originalBuildStatuses.has(nodeId)) { + const currentStatus = useFlowStore.getState().flowBuildStatus[nodeId]; + originalBuildStatuses.set( + nodeId, + currentStatus ? { ...currentStatus } : undefined, + ); + } flowStore.updateBuildStatus([nodeId], buildStatus); nodeIds.add(nodeId); @@ -133,9 +162,19 @@ export function applyStateDelta( // bridge's final ``finish()`` clears edges, and they never animate during // the run. if (value.status === "running") { - flowStore.updateEdgesRunningByNodes([nodeId], true); + if (runningNodeIds) { + runningNodeIds.add(nodeId); + flowStore.clearAndSetEdgesRunning([...runningNodeIds]); + } else { + flowStore.updateEdgesRunningByNodes([nodeId], true); + } } else if (value.status === "success" || value.status === "error") { - flowStore.updateEdgesRunningByNodes([nodeId], false); + if (runningNodeIds) { + runningNodeIds.delete(nodeId); + flowStore.clearAndSetEdgesRunning([...runningNodeIds]); + } else { + flowStore.updateEdgesRunningByNodes([nodeId], false); + } } // Only the final per-vertex emission carries the result data; running @@ -203,6 +242,11 @@ export async function runFlowAGUI( const flowStore = useFlowStore.getState(); const setErrorData = useAlertStore.getState().setErrorData; const touchedNodeIds = new Set(); + const runningNodeIds = new Set(); + const originalBuildStatuses = new Map< + string, + FlowBuildStatusEntry | undefined + >(); // Server derives run_id from session_id/flow_id and announces it via // RUN_STARTED. Until that arrives we have no id to stamp on flow-pool // entries, so we fall back to an empty string (matches the legacy @@ -214,6 +258,91 @@ export async function runFlowAGUI( // failure. Without this, a silent ``complete:`` would leave ``buildInfo`` // ``null`` and ``trackFlowBuild`` would mis-record the run as success. let terminalEventSeen = false; + let markRunningAsError = false; + + const persistBuildDuration = (data: unknown) => { + const buildDuration = (data as { build_duration?: number } | null) + ?.build_duration; + if (buildDuration == null) return; + + const durationMs = buildDuration * 1000; + const currentFlowStore = useFlowStore.getState(); + currentFlowStore.setBuildDuration(durationMs); + + const found = findLastBotMessage( + currentFlowStore.buildingFlowId ?? undefined, + currentFlowStore.buildingSessionId ?? undefined, + ); + if (!found) return; + if (found.message.properties?.build_duration != null) return; + + useMessagesStore.getState().updateMessagePartial({ + id: found.message.id, + properties: { + ...found.message.properties, + build_duration: durationMs, + }, + }); + + updateMessageProperties(found.message.id!, found.queryKey, { + build_duration: durationMs, + }); + persistMessageProperties(found.message.id!, { + ...found.message, + properties: { + ...found.message.properties, + build_duration: durationMs, + }, + }); + }; + + const appendLogEvent = (data: unknown) => { + const { component_id, output, name, message, type: logType } = + (data ?? {}) as { + component_id?: string; + output?: string; + } & Partial; + if (!component_id || !output) { + console.error( + "[runFlowAGUI] Received malformed log event; missing component_id or output", + data, + ); + return; + } + useFlowStore.getState().appendLogToFlowPool(component_id, output, { + name, + message, + type: logType, + } as LogsLogType); + }; + + const markRunningNodesFailed = () => { + const current = useFlowStore.getState(); + const flowBuildStatus = { ...current.flowBuildStatus }; + for (const nodeId of touchedNodeIds) { + if (flowBuildStatus[nodeId]?.status === BuildStatus.BUILDING) { + flowBuildStatus[nodeId] = { + ...flowBuildStatus[nodeId], + status: BuildStatus.ERROR, + }; + } + } + useFlowStore.setState({ flowBuildStatus }); + }; + + const restoreOriginalBuildStatuses = () => { + const current = useFlowStore.getState(); + const flowBuildStatus = { ...current.flowBuildStatus }; + for (const [nodeId, status] of originalBuildStatuses) { + if (flowBuildStatus[nodeId]?.status !== BuildStatus.BUILDING) continue; + if (status) { + flowBuildStatus[nodeId] = { ...status }; + } else { + delete flowBuildStatus[nodeId]; + } + } + useFlowStore.setState({ flowBuildStatus }); + }; // `agent.run` still needs a `RunAgentInput` for the client-side apply // pipeline (subscriber correlation); the actual wire body is the native @@ -237,14 +366,26 @@ export async function runFlowAGUI( setRunId: (r) => { runId = r; }, - applyDelta: (ops) => applyStateDelta(ops, runId, touchedNodeIds), + applyDelta: (ops) => + applyStateDelta( + ops, + runId, + touchedNodeIds, + runningNodeIds, + originalBuildStatuses, + ), handleCustomEvent: (eventType, data) => handleMessageEvent(eventType, data), + handleEndEvent: persistBuildDuration, + handleLogEvent: appendLogEvent, onFinished: () => { terminalEventSeen = true; - flowStore.setBuildInfo({ success: true }); + if (!opts.silent) { + flowStore.setBuildInfo({ success: true }); + } }, onError: (message) => { terminalEventSeen = true; + markRunningAsError = true; flowStore.setBuildInfo({ error: [message], success: false }); setErrorData({ title: "Workflow run failed", list: [message] }); }, @@ -255,9 +396,15 @@ export async function runFlowAGUI( const finish = () => { if (settled) return; settled = true; - flowStore.updateEdgesRunningByNodes([...touchedNodeIds], false); + flowStore.clearAndSetEdgesRunning([]); flowStore.setIsBuilding(false); - flowStore.revertBuiltStatusFromBuilding(); + if (markRunningAsError) { + markRunningNodesFailed(); + } else if (!terminalEventSeen) { + restoreOriginalBuildStatuses(); + } else { + flowStore.revertBuiltStatusFromBuilding(); + } resolve(); }; @@ -299,6 +446,7 @@ export async function runFlowAGUI( } }, error: (err: Error) => { + markRunningAsError = true; flowStore.setBuildInfo({ error: [err.message], success: false }); setErrorData({ title: "Workflow run failed", list: [err.message] }); subscription.unsubscribe(); @@ -310,6 +458,7 @@ export async function runFlowAGUI( // truncated response, proxy timeout). Record an error so analytics // doesn't treat a silent close as success. if (!terminalEventSeen) { + markRunningAsError = true; flowStore.setBuildInfo({ error: ["Workflow run ended unexpectedly"], success: false, diff --git a/src/frontend/src/stores/__tests__/flowStore.test.ts b/src/frontend/src/stores/__tests__/flowStore.test.ts index 809ec66545..5040ee808b 100644 --- a/src/frontend/src/stores/__tests__/flowStore.test.ts +++ b/src/frontend/src/stores/__tests__/flowStore.test.ts @@ -1,4 +1,5 @@ import { act, renderHook } from "@testing-library/react"; +import { v5 as uuidv5 } from "uuid"; // Mock all the complex dependencies jest.mock("@xyflow/react", () => ({ @@ -124,6 +125,7 @@ jest.mock("../../utils/reactflowUtils", () => { import { checkCodeValidity } from "@/CustomNodes/helpers/check-code-validity"; import type { LogsLogType, VertexBuildTypeAPI } from "@/types/api"; import type { AllNodeType, EdgeType } from "@/types/flow"; +import useAuthStore from "../authStore"; import useFlowStore, { completeNodeUpdate, recomputeComponentsToUpdateIfNeeded, @@ -1170,6 +1172,12 @@ describe("useFlowStore", () => { beforeEach(() => { mockedRunFlow.mockReset(); trackFlowBuildMock.mockReset(); + useAuthStore.setState({ + isAuthenticated: false, + autoLogin: null, + userData: null, + }); + useUtilityStore.setState({ clientId: "" }); act(() => { useFlowStore.setState({ nodes: [], @@ -1212,5 +1220,59 @@ describe("useFlowStore", () => { error: errorList, }); }); + + it("passes silent and sets the active building session before the AG-UI run", async () => { + mockedRunFlow.mockImplementation(async () => { + const state = useFlowStore.getState(); + expect(state.buildingFlowId).toBe("flow-abc"); + expect(state.buildingSessionId).toBe("session-123"); + useFlowStore.setState({ buildInfo: null }); + }); + + await useFlowStore.getState().buildFlow({ + silent: true, + session: "session-123", + }); + + expect(mockedRunFlow).toHaveBeenCalledWith( + expect.objectContaining({ + flowId: "flow-abc", + threadId: "session-123", + silent: true, + }), + ); + expect(trackFlowBuildMock).toHaveBeenCalledWith("Test Flow", false, { + flowId: "flow-abc", + }); + }); + + it("uses the public playground virtual flow id for the active building session", async () => { + const expectedFlowId = uuidv5("client-123_flow-abc", uuidv5.DNS); + + useUtilityStore.setState({ clientId: "client-123" }); + useAuthStore.setState({ + isAuthenticated: false, + autoLogin: true, + userData: null, + }); + useFlowStore.setState({ playgroundPage: true }); + mockedRunFlow.mockImplementation(async () => { + const state = useFlowStore.getState(); + expect(state.buildingFlowId).toBe(expectedFlowId); + expect(state.buildingSessionId).toBe("session-123"); + useFlowStore.setState({ buildInfo: null }); + }); + + await useFlowStore.getState().buildFlow({ + session: "session-123", + }); + + expect(mockedRunFlow).toHaveBeenCalledWith( + expect.objectContaining({ + flowId: "flow-abc", + threadId: "session-123", + }), + ); + }); }); }); diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts index 4e82bf5bbd..9c9db0d4d5 100644 --- a/src/frontend/src/stores/flowStore.ts +++ b/src/frontend/src/stores/flowStore.ts @@ -7,6 +7,7 @@ import { type NodeChange, } from "@xyflow/react"; import { cloneDeep } from "lodash"; +import { v5 as uuidv5 } from "uuid"; import { create } from "zustand"; import { checkCodeValidity } from "@/CustomNodes/helpers/check-code-validity"; import { queryClient } from "@/contexts"; @@ -46,6 +47,7 @@ import { } from "../utils/reactflowUtils"; import { getInputsAndOutputs } from "../utils/storeUtils"; import useAlertStore from "./alertStore"; +import useAuthStore from "./authStore"; import { useDarkStore } from "./darkStore"; import useFlowsManagerStore from "./flowsManagerStore"; import { useGlobalVariablesStore } from "./globalVariablesStore/globalVariables"; @@ -929,6 +931,18 @@ const useFlowStore = create((set, get) => ({ // Always run through the v2 workflows endpoint. Current frontend nodes // + edges are sent so unsaved tweaks (dropdowns, text inputs) run as // the user sees them. + let buildingFlowId = currentFlow!.id; + if (get().playgroundPage) { + const authState = useAuthStore.getState(); + const visitorId = + authState.isAuthenticated && + authState.autoLogin === false && + authState.userData?.id + ? authState.userData.id + : useUtilityStore.getState().clientId; + buildingFlowId = uuidv5(`${visitorId}_${currentFlow!.id}`, uuidv5.DNS); + } + get().setBuildingSession(buildingFlowId, session ?? null); await runFlowAGUI({ flowId: currentFlow!.id, message: input_value, @@ -938,6 +952,7 @@ const useFlowStore = create((set, get) => ({ flowData: { nodes: get().nodes, edges: get().edges }, files, signal: buildController.signal, + silent, }); // Invalidate KB-related caches so any KnowledgeIngestion node that ran @@ -950,8 +965,9 @@ const useFlowStore = create((set, get) => ({ // Mirror the v1 build callbacks' analytics: every actual build attempt // logs a flow-build event with success/error and (on failure) the error - // list. `runFlowAGUI` always resolves and writes the outcome into - // `buildInfo`, so reading it back is the cleanest success/error signal. + // list. `runFlowAGUI` always resolves and writes failures into + // `buildInfo`; silent successful runs leave it null and are tracked as + // success. const finalBuildInfo = get().buildInfo; const hasError = finalBuildInfo?.success === false; trackFlowBuild(currentFlow?.name ?? "Unknown", hasError, { From c84e965c66f675841a4928f577dcbff808c73a3a Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 14:47:51 -0300 Subject: [PATCH 04/11] fix(api/v2): signal cross-worker workflow stops LE-1389 --- src/backend/base/langflow/api/v2/workflow.py | 20 +- .../tests/unit/api/v2/test_workflow.py | 21 ++- .../tests/unit/api/v2/test_workflow_agui.py | 172 ++++++++++++++++++ 3 files changed, 202 insertions(+), 11 deletions(-) diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index f7a04b9e28..da6e8a8ba4 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -102,6 +102,13 @@ async def generate_flow_events(*args, **kwargs) -> None: await _generate_flow_events(*args, **kwargs) +async def _cancel_workflow_queue_job(job_id: str) -> bool: + """Lazily call the shared build cancellation path to avoid import cycles.""" + from langflow.api.build import cancel_flow_build + + return await cancel_flow_build(job_id=job_id, queue_service=get_queue_service()) + + def _unknown_protocol_http_exception(exc: UnknownStreamProtocolError) -> HTTPException: """Build the 422 response shared by ``stream`` and ``background`` paths. @@ -1373,7 +1380,7 @@ async def stop_workflow( try: try: - await get_queue_service().cancel_job(str(job_id)) + cancelled = await _cancel_workflow_queue_job(str(job_id)) except asyncio.CancelledError as exc: message_code = exc.args[0] if exc.args else "UNKNOWN" if message_code != "LANGFLOW_USER_CANCELLED": @@ -1386,6 +1393,17 @@ async def stop_workflow( "job_id": str(job_id), }, ) from exc + cancelled = True + if not cancelled: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={ + "error": "Cancellation unavailable", + "code": "WORKFLOW_CANCEL_UNAVAILABLE", + "message": f"Unable to confirm cancellation for job {job_id}", + "job_id": str(job_id), + }, + ) await job_service.update_job_status(job_id, JobStatus.CANCELLED) # Release the in-memory buffer and wake any re-attach waiters so they # see a clean stream-end instead of hanging on ``_cond.wait()``. diff --git a/src/backend/tests/unit/api/v2/test_workflow.py b/src/backend/tests/unit/api/v2/test_workflow.py index 6672e7201d..519b7d6e85 100644 --- a/src/backend/tests/unit/api/v2/test_workflow.py +++ b/src/backend/tests/unit/api/v2/test_workflow.py @@ -281,16 +281,14 @@ class TestWorkflowStop: with ( patch("langflow.api.v2.workflow.get_job_service") as mock_get_job_service, - patch("langflow.api.v2.workflow.get_task_service") as mock_get_task_service, + patch("langflow.api.v2.workflow._cancel_workflow_queue_job") as mock_cancel_workflow_queue_job, ): mock_job_service = MagicMock() mock_job_service.get_job_by_job_id = AsyncMock(return_value=mock_job) mock_job_service.update_job_status = AsyncMock() mock_get_job_service.return_value = mock_job_service - mock_task_service = MagicMock() - mock_task_service.revoke_task = AsyncMock(return_value=True) - mock_get_task_service.return_value = mock_task_service + mock_cancel_workflow_queue_job.return_value = True headers = {"x-api-key": created_api_key.api_key} response = await client.post("api/v2/workflows/stop", json={"job_id": job_id}, headers=headers) @@ -299,7 +297,7 @@ class TestWorkflowStop: result = response.json() assert result["job_id"] == job_id assert "cancelled successfully" in result["message"] - mock_task_service.revoke_task.assert_called_once() + mock_cancel_workflow_queue_job.assert_awaited_once_with(job_id) mock_job_service.update_job_status.assert_called_once_with(UUID(job_id), JobStatus.CANCELLED) async def test_stop_workflow_not_found( @@ -574,15 +572,18 @@ class TestWorkflowIDORProtection: try: headers = {"x-api-key": created_api_key.api_key} - response = await client.post( - "api/v2/workflows/stop", - json={"job_id": str(job_id)}, - headers=headers, - ) + with patch("langflow.api.v2.workflow._cancel_workflow_queue_job") as mock_cancel_workflow_queue_job: + mock_cancel_workflow_queue_job.return_value = True + response = await client.post( + "api/v2/workflows/stop", + json={"job_id": str(job_id)}, + headers=headers, + ) assert response.status_code == 200, ( "Legacy jobs with user_id=None must not be blocked by the ownership check" ) + mock_cancel_workflow_queue_job.assert_awaited_once_with(str(job_id)) finally: async with session_scope() as session: db_job = await session.get(Job, job_id) diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index f9f6d60c1b..d58aab0b5a 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -227,6 +227,34 @@ class TestAGUIModeDispatch: result = response.json() assert result["status"] == "completed" + async def test_sync_mode_rejects_live_canvas_only_fields( + self, + client: AsyncClient, + created_api_key, + empty_flow, + ): + """Sync runs must reject fields that only stream/background executes.""" + body = _agui_body(empty_flow, mode="sync") + body.update( + { + "data": {"nodes": [], "edges": []}, + "files": ["tmp/upload.txt"], + "start_component_id": "input-1", + "stop_component_id": "output-1", + } + ) + + response = await client.post( + "api/v2/workflows", + json=body, + headers={"x-api-key": created_api_key.api_key}, + ) + + assert response.status_code == 422 + detail = response.json()["detail"] + assert detail["code"] == "SYNC_MODE_UNSUPPORTED_FIELDS" + assert detail["fields"] == ["data", "files", "start_component_id", "stop_component_id"] + class TestV2WorkflowAdmission: """Route-level admission checks before workflow execution dispatch.""" @@ -301,6 +329,41 @@ class TestV2WorkflowAdmission: assert exc_info.value.status_code == 404 assert exc_info.value.detail["code"] == "FLOW_NOT_FOUND" + async def test_private_route_applies_component_policy_gate(self, monkeypatch: pytest.MonkeyPatch): + """The authenticated v2 route must run server-side component policy validation.""" + from langflow.api.v2 import workflow as workflow_module + from lfx.schema.workflow import WorkflowRunRequest + from lfx.utils.flow_validation import CustomComponentValidationError + + flow_id = uuid4() + flow = SimpleNamespace( + id=flow_id, + user_id=uuid4(), + workspace_id=None, + folder_id=None, + data={"nodes": [], "edges": []}, + name="private", + ) + + def _reject(_flow_data): + message = "custom components are disabled" + raise CustomComponentValidationError(message) + + monkeypatch.setattr(workflow_module, "get_flow_by_id_or_endpoint_name", AsyncMock(return_value=flow)) + monkeypatch.setattr(workflow_module, "ensure_flow_permission", AsyncMock(return_value=None)) + monkeypatch.setattr(workflow_module, "validate_flow_for_current_settings", _reject) + + with pytest.raises(HTTPException) as exc_info: + await workflow_module.execute_workflow( + WorkflowRunRequest(flow_id=str(flow_id), input_value="hi", mode="stream"), + background_tasks=SimpleNamespace(), + http_request=SimpleNamespace(), + current_user=SimpleNamespace(id=flow.user_id), + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "custom components are disabled" + class TestAGUIStreaming: """mode=stream returns an AG-UI server-sent event stream.""" @@ -365,6 +428,57 @@ class TestAGUIStreaming: assert seen_maxsize == [2] assert event_types == ["token", "token", "error"] + async def test_agui_stream_emits_end_side_channel_for_build_duration(self, monkeypatch: pytest.MonkeyPatch): + """The AG-UI stream must preserve v1 end payloads for chat build-duration persistence.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + async def fake_generate_flow_events(**kwargs): + event_queue = kwargs["event_manager"].queue + payload = json.dumps({"event": "end", "data": {"build_duration": 1.25}}).encode() + event_queue.put_nowait(("end-1", payload, time.time())) + await event_queue.put((None, None, time.time())) + + class FakeAdapter: + name = "agui" + + def initial_events(self): + return [] + + def final_events(self): + return [] + + def translate(self, _event_type, _event_data): + return [] + + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + frames = [ + frame + async for frame, _event_type in workflow_module._stream_event_frames( + adapter=FakeAdapter(), + flow_id=uuid4(), + flow_name="flow", + background_tasks=SimpleNamespace(add_task=lambda *_args, **_kwargs: None), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="", mode="stream"), + current_user=SimpleNamespace(id=uuid4()), + ) + ] + + custom_events = [ + json.loads(line.removeprefix("data:").strip()) + for frame in frames + for line in frame.decode("utf-8").splitlines() + if line.startswith("data:") + ] + assert custom_events == [ + { + "type": "CUSTOM", + "name": "langflow.event", + "value": {"event_type": "end", "data": {"build_duration": 1.25}}, + } + ] + async def test_stream_emits_run_lifecycle_events( self, client: AsyncClient, @@ -1901,3 +2015,61 @@ class TestStopWorkflowEndToEnd: with contextlib.suppress(BaseException): await get_queue_service().cleanup_job(job_id) await workflow_module._clear_background_run(job_id) + + async def test_stop_signals_cross_worker_queue_owner_before_marking_cancelled( + self, + monkeypatch: pytest.MonkeyPatch, + ): + """Redis-backed jobs owned by another worker must be cancelled by pub/sub signal.""" + from langflow.api.v2 import workflow as workflow_module + + job_id = uuid4() + current_user_id = uuid4() + + class FakeJobService: + def __init__(self) -> None: + self.updates: list[tuple[object, object]] = [] + + async def get_job_by_job_id(self, seen_job_id, user_id=None): + assert seen_job_id == job_id + assert user_id == current_user_id + return SimpleNamespace( + type=workflow_module.JobType.WORKFLOW, + status=workflow_module.JobStatus.IN_PROGRESS, + ) + + async def update_job_status(self, seen_job_id, status): + self.updates.append((seen_job_id, status)) + + class CrossWorkerQueueService: + cross_worker_cancel_enabled = True + + def __init__(self) -> None: + self.local_cancels: list[str] = [] + self.signals: list[str] = [] + + def get_queue_data(self, seen_job_id): + assert seen_job_id == str(job_id) + return SimpleNamespace(), SimpleNamespace(), None, None + + async def cancel_job(self, seen_job_id): + self.local_cancels.append(seen_job_id) + + async def signal_cancel(self, seen_job_id): + self.signals.append(seen_job_id) + return 0 + + job_service = FakeJobService() + queue_service = CrossWorkerQueueService() + monkeypatch.setattr(workflow_module, "get_job_service", lambda: job_service) + monkeypatch.setattr(workflow_module, "get_queue_service", lambda: queue_service) + + response = await workflow_module.stop_workflow( + workflow_module.WorkflowStopRequest(job_id=job_id), + current_user=SimpleNamespace(id=current_user_id), + ) + + assert str(response.job_id) == str(job_id) + assert queue_service.local_cancels == [] + assert queue_service.signals == [str(job_id)] + assert job_service.updates == [(job_id, workflow_module.JobStatus.CANCELLED)] From 54fc55082fc3c9a20171fa79620b8ea7a52ab0a1 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 15:00:33 -0300 Subject: [PATCH 05/11] fix(api/v2): report unconfirmed workflow stops LE-1389 --- src/backend/base/langflow/api/v2/workflow.py | 6 ++- .../tests/unit/api/v2/test_workflow.py | 39 +++++++++++++++++++ .../tests/unit/api/v2/test_workflow_agui.py | 7 ++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index da6e8a8ba4..2b4e67cb1a 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -105,8 +105,12 @@ async def generate_flow_events(*args, **kwargs) -> None: async def _cancel_workflow_queue_job(job_id: str) -> bool: """Lazily call the shared build cancellation path to avoid import cycles.""" from langflow.api.build import cancel_flow_build + from langflow.services.job_queue.service import JobQueueNotFoundError - return await cancel_flow_build(job_id=job_id, queue_service=get_queue_service()) + try: + return await cancel_flow_build(job_id=job_id, queue_service=get_queue_service()) + except JobQueueNotFoundError: + return False def _unknown_protocol_http_exception(exc: UnknownStreamProtocolError) -> HTTPException: diff --git a/src/backend/tests/unit/api/v2/test_workflow.py b/src/backend/tests/unit/api/v2/test_workflow.py index 519b7d6e85..3f365865c5 100644 --- a/src/backend/tests/unit/api/v2/test_workflow.py +++ b/src/backend/tests/unit/api/v2/test_workflow.py @@ -300,6 +300,45 @@ class TestWorkflowStop: mock_cancel_workflow_queue_job.assert_awaited_once_with(job_id) mock_job_service.update_job_status.assert_called_once_with(UUID(job_id), JobStatus.CANCELLED) + async def test_stop_workflow_returns_503_when_queue_cancel_cannot_be_confirmed( + self, + client: AsyncClient, + created_api_key, + ): + """Do not mark persisted jobs cancelled when no queue owner can be reached.""" + from langflow.services.job_queue.service import JobQueueNotFoundError + + job_id = str(uuid4()) + + mock_job = MagicMock() + mock_job.job_id = job_id + mock_job.status = JobStatus.IN_PROGRESS + mock_job.type = JobType.WORKFLOW + mock_job.user_id = None + + class MissingQueueService: + def get_queue_data(self, seen_job_id): + raise JobQueueNotFoundError(seen_job_id) + + with ( + patch("langflow.api.v2.workflow.get_job_service") as mock_get_job_service, + patch("langflow.api.v2.workflow.get_queue_service", return_value=MissingQueueService()), + ): + mock_job_service = MagicMock() + mock_job_service.get_job_by_job_id = AsyncMock(return_value=mock_job) + mock_job_service.update_job_status = AsyncMock() + mock_get_job_service.return_value = mock_job_service + + response = await client.post( + "api/v2/workflows/stop", + json={"job_id": job_id}, + headers={"x-api-key": created_api_key.api_key}, + ) + + assert response.status_code == 503 + assert response.json()["detail"]["code"] == "WORKFLOW_CANCEL_UNAVAILABLE" + mock_job_service.update_job_status.assert_not_awaited() + async def test_stop_workflow_not_found( self, client: AsyncClient, diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index d58aab0b5a..a8fe06e5b1 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -2025,6 +2025,7 @@ class TestStopWorkflowEndToEnd: job_id = uuid4() current_user_id = uuid4() + events: list[tuple[str, object, object | None]] = [] class FakeJobService: def __init__(self) -> None: @@ -2039,6 +2040,7 @@ class TestStopWorkflowEndToEnd: ) async def update_job_status(self, seen_job_id, status): + events.append(("update", seen_job_id, status)) self.updates.append((seen_job_id, status)) class CrossWorkerQueueService: @@ -2056,6 +2058,7 @@ class TestStopWorkflowEndToEnd: self.local_cancels.append(seen_job_id) async def signal_cancel(self, seen_job_id): + events.append(("signal", seen_job_id, None)) self.signals.append(seen_job_id) return 0 @@ -2073,3 +2076,7 @@ class TestStopWorkflowEndToEnd: assert queue_service.local_cancels == [] assert queue_service.signals == [str(job_id)] assert job_service.updates == [(job_id, workflow_module.JobStatus.CANCELLED)] + assert events == [ + ("signal", str(job_id), None), + ("update", job_id, workflow_module.JobStatus.CANCELLED), + ] From 60469a698ee425ab6c6bd9387b0f3ae0a86534f0 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 15:23:41 -0300 Subject: [PATCH 06/11] fix(api/v2): keep background workflows out of polling watchdog LE-1389 --- src/backend/base/langflow/api/v2/workflow.py | 1 - .../tests/unit/api/v2/test_workflow_agui.py | 54 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index 2b4e67cb1a..43e4b87d52 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -1123,7 +1123,6 @@ async def execute_workflow_background( try: queue_service = get_queue_service() queue_service.create_queue(job_id_str) - await queue_service.register_job_owner(job_id_str, current_user.id) queue_service.start_job( job_id_str, _buffer_background_run( diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index a8fe06e5b1..7dd04af556 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -1904,6 +1904,60 @@ class TestBufferBackgroundRunUnknownProtocolGuard: _REGISTRY.update(snapshot) +class TestExecuteWorkflowBackgroundQueueOwnership: + async def test_background_jobs_do_not_register_polling_owner(self, monkeypatch: pytest.MonkeyPatch): + """Background jobs must not opt into Redis' polling-owner watchdog. + + ``register_job_owner`` is the Redis queue service's signal that a job + has an active polling/streaming client expected to refresh + ``touch_activity``. v2 background jobs run unattended and are stopped + via the queue cancel side-channel plus the persisted workflow Job row, + so registering them would let the polling watchdog reclaim long runs. + """ + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + job_id = uuid4() + current_user_id = uuid4() + flow_id = uuid4() + created_jobs: list[object] = [] + registered_owners: list[tuple[str, object]] = [] + started_jobs: list[str] = [] + + class FakeJobService: + async def create_job(self, **kwargs): + created_jobs.append(kwargs["job_id"]) + + class FakeQueueService: + def create_queue(self, _seen_job_id): + return SimpleNamespace(), SimpleNamespace() + + async def register_job_owner(self, seen_job_id, user_id): + registered_owners.append((seen_job_id, user_id)) + + def start_job(self, seen_job_id, task_coro): + started_jobs.append(seen_job_id) + task_coro.close() + + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + monkeypatch.setattr(workflow_module, "get_queue_service", lambda: FakeQueueService()) + monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {}) + + response = await workflow_module.execute_workflow_background( + parsed=ParsedWorkflowRun(flow_id=str(flow_id), mode="background"), + flow=SimpleNamespace(id=flow_id, name="Background Flow"), + job_id=job_id, + current_user=SimpleNamespace(id=current_user_id), + http_request=SimpleNamespace(), + stream_protocol="agui", + ) + + assert str(response.job_id) == str(job_id) + assert created_jobs == [job_id] + assert started_jobs == [str(job_id)] + assert registered_owners == [] + + class TestStopWorkflowEndToEnd: """The full ``POST /workflows/stop`` HTTP flow clears the in-memory buffer too. From 38dc8995e993bc63bf567809b2766f4c37ac82e7 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Thu, 11 Jun 2026 11:11:40 -0300 Subject: [PATCH 07/11] Fix AG-UI workflow lifecycle edges --- .../base/langflow/api/v2/adapters/__init__.py | 3 + .../base/langflow/api/v2/adapters/agui.py | 9 +- .../base/langflow/api/v2/adapters/langflow.py | 4 + .../base/langflow/api/v2/agui_translator.py | 68 ++-- src/backend/base/langflow/api/v2/workflow.py | 126 ++++-- .../unit/api/v2/adapters/test_agui_adapter.py | 23 ++ .../api/v2/adapters/test_langflow_adapter.py | 7 + .../unit/api/v2/adapters/test_registry.py | 3 + .../tests/unit/api/v2/test_agui_translator.py | 76 ++-- .../tests/unit/api/v2/test_workflow_agui.py | 371 +++++++++++++++++- .../lfx/custom/custom_component/component.py | 3 +- src/lfx/src/lfx/schema/workflow.py | 1 + .../custom/custom_component/test_component.py | 25 ++ 13 files changed, 608 insertions(+), 111 deletions(-) diff --git a/src/backend/base/langflow/api/v2/adapters/__init__.py b/src/backend/base/langflow/api/v2/adapters/__init__.py index 3ffe9e18e2..75c9006300 100644 --- a/src/backend/base/langflow/api/v2/adapters/__init__.py +++ b/src/backend/base/langflow/api/v2/adapters/__init__.py @@ -68,6 +68,9 @@ class StreamAdapter(Protocol): def error_events(self, error: BaseException) -> Iterable[StreamEvent]: """Events to emit when the run errors mid-stream.""" + def cancel_events(self, reason: str) -> Iterable[StreamEvent]: + """Events to emit when the run is cancelled by the user.""" + @property def terminal_error_type(self) -> str | None: """Event-type the buffer task watches to flip a background job to FAILED. diff --git a/src/backend/base/langflow/api/v2/adapters/agui.py b/src/backend/base/langflow/api/v2/adapters/agui.py index 495e46cfdd..ac6981cb6a 100644 --- a/src/backend/base/langflow/api/v2/adapters/agui.py +++ b/src/backend/base/langflow/api/v2/adapters/agui.py @@ -10,7 +10,7 @@ from __future__ import annotations from collections.abc import Iterable from typing import Any, ClassVar -from ag_ui.core import BaseEvent, RunErrorEvent +from ag_ui.core import BaseEvent from langflow.api.v2.adapters import ( StreamAdapterContext, @@ -56,7 +56,12 @@ class AGUIAdapter: def error_events(self, error: BaseException) -> Iterable[StreamEvent]: # Fallback path: emitted by the dispatcher when on_error itself fails # or when no error event has reached the translator. - return [_to_stream_event(RunErrorEvent(message=str(error)))] + return [_to_stream_event(e) for e in self._translator.translate("error", {"error": str(error)})] + + def cancel_events(self, reason: str) -> Iterable[StreamEvent]: + # AG-UI has no cancellation primitive in the local event model; route + # through the translator so any open text lifecycle is closed first. + return [_to_stream_event(e) for e in self._translator.translate("error", {"error": reason})] @property def terminal_error_type(self) -> str | None: diff --git a/src/backend/base/langflow/api/v2/adapters/langflow.py b/src/backend/base/langflow/api/v2/adapters/langflow.py index cdc3ec250d..043e225f35 100644 --- a/src/backend/base/langflow/api/v2/adapters/langflow.py +++ b/src/backend/base/langflow/api/v2/adapters/langflow.py @@ -87,6 +87,10 @@ class LangflowAdapter: payload = {"event": "error", "data": {"error": str(error)}} return (StreamEvent(type="error", data_json=json.dumps(payload)),) + def cancel_events(self, reason: str) -> Iterable[StreamEvent]: + payload = {"event": "error", "data": {"error": reason}} + return (StreamEvent(type="error", data_json=json.dumps(payload)),) + @property def terminal_error_type(self) -> str | None: return "error" diff --git a/src/backend/base/langflow/api/v2/agui_translator.py b/src/backend/base/langflow/api/v2/agui_translator.py index 727e71e56e..f01ff6614a 100644 --- a/src/backend/base/langflow/api/v2/agui_translator.py +++ b/src/backend/base/langflow/api/v2/agui_translator.py @@ -47,9 +47,10 @@ class AGUITranslator: def __init__(self, run_id: str, thread_id: str) -> None: self.run_id = run_id self.thread_id = thread_id - # Id of the text message currently being streamed by ``token`` events, - # or ``None`` when no message is open. - self._open_message_id: str | None = None + # Text messages currently being streamed by ``token`` events. Langflow + # can run vertices in parallel, so token streams may interleave; AG-UI + # consumers correlate content by message_id. + self._open_message_ids: dict[str, None] = {} # Message ids already emitted as a complete (non-streamed) text message. self._emitted_text_message_ids: set[str] = set() # Tool-call ids already emitted as TOOL_CALL_START / already resolved @@ -95,11 +96,11 @@ class AGUITranslator: # streamed message and must stay transparent, or the message would be # split into multiple START/END pairs reusing an already-ended id. if event_type == "end": - events = self._close_open_message() + events = self._close_open_messages() events.append(RunFinishedEvent(run_id=self.run_id, thread_id=self.thread_id)) return events if event_type == "error": - events = self._close_open_message() + events = self._close_open_messages() # The ``error`` payload varies by emission path: a full ErrorMessage # dump carries the reason in ``text``; the minimal path sends # ``{"error": str}``. @@ -112,9 +113,9 @@ class AGUITranslator: """Map a ``token`` event to text-message events. The first token of a message opens it with ``TEXT_MESSAGE_START``; a - token for a different message id closes the previous one first. A - token whose id was already ended via ``add_message`` (or a prior - token boundary) is dropped: re-opening it would emit a second + token for a different message id opens another message without closing + existing streams. A token whose id was already ended via + ``add_message`` is dropped: re-opening it would emit a second ``TEXT_MESSAGE_START`` for an id the protocol considers closed. """ message_id = str(data.get("id") or "") @@ -123,14 +124,13 @@ class AGUITranslator: # be correlated. Dropping the event is preferable to emitting a # malformed stream with empty message_ids. return [] - if message_id in self._emitted_text_message_ids and self._open_message_id != message_id: + if message_id in self._emitted_text_message_ids and message_id not in self._open_message_ids: return [] chunk = data.get("chunk", "") events: list[BaseEvent] = [] - if self._open_message_id != message_id: - events.extend(self._close_open_message()) + if message_id not in self._open_message_ids: events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) - self._open_message_id = message_id + self._open_message_ids[message_id] = None events.append(TextMessageContentEvent(message_id=message_id, delta=chunk)) return events @@ -197,18 +197,20 @@ class AGUITranslator: ) # Message text. - if message_id and message_id == self._open_message_id: + message_state = self._message_state(data) + if message_id and message_id in self._open_message_ids: + if message_state == "partial": + return events # Finalizer of a token-streamed message: close it. The text was # already streamed token by token, so it must not be re-emitted now # or by any later add_message that re-fires for the same id. - self._emitted_text_message_ids.add(message_id) - events.extend(self._close_open_message()) + events.extend(self._close_message(message_id)) else: text = data.get("text") or "" # Skip text-message lifecycle emission without a stable message_id; # tool-call events above are namespaced by block/content index so # they can still ride a missing id, but TEXT_MESSAGE_* cannot. - if text and message_id and message_id not in self._emitted_text_message_ids: + if text and message_id and message_state != "partial" and message_id not in self._emitted_text_message_ids: self._emitted_text_message_ids.add(message_id) events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) events.append(TextMessageContentEvent(message_id=message_id, delta=text)) @@ -275,6 +277,14 @@ class AGUITranslator: ) ] + @staticmethod + def _message_state(data: dict) -> str | None: + properties = data.get("properties") + if not isinstance(properties, dict): + return None + state = properties.get("state") + return str(state) if state is not None else None + @staticmethod def _set_node(node_id: str, status: str, output: object) -> dict: """Build the RFC 6902 op that writes a node's state. @@ -285,18 +295,18 @@ class AGUITranslator: """ return {"op": "add", "path": f"/nodes/{node_id}", "value": {"status": status, "output": output}} - def _close_open_message(self) -> list[BaseEvent]: - """Emit ``TEXT_MESSAGE_END`` for the open message, if any. - - The closed id is recorded in ``_emitted_text_message_ids`` so a later - token (e.g. an interleaved ``A, B, A`` sequence) cannot re-open it - and emit a second ``TEXT_MESSAGE_START`` for an id the protocol - already considers closed. - """ - if self._open_message_id is None: + def _close_message(self, message_id: str) -> list[BaseEvent]: + """Emit ``TEXT_MESSAGE_END`` for one open message, if present.""" + if message_id not in self._open_message_ids: return [] - closed_id = self._open_message_id - end = TextMessageEndEvent(message_id=closed_id) - self._emitted_text_message_ids.add(closed_id) - self._open_message_id = None + self._open_message_ids.pop(message_id, None) + end = TextMessageEndEvent(message_id=message_id) + self._emitted_text_message_ids.add(message_id) return [end] + + def _close_open_messages(self) -> list[BaseEvent]: + """Emit ``TEXT_MESSAGE_END`` for every open message.""" + events: list[BaseEvent] = [] + for message_id in list(self._open_message_ids): + events.extend(self._close_message(message_id)) + return events diff --git a/src/backend/base/langflow/api/v2/workflow.py b/src/backend/base/langflow/api/v2/workflow.py index 43e4b87d52..2bbbf4f2ec 100644 --- a/src/backend/base/langflow/api/v2/workflow.py +++ b/src/backend/base/langflow/api/v2/workflow.py @@ -24,7 +24,7 @@ import asyncio import contextlib import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable, Iterable from copy import deepcopy from typing import Annotated from uuid import UUID, uuid4 @@ -736,11 +736,10 @@ async def _stream_event_frames( # RUN_ERROR for AG-UI, ``error`` for langflow) so the buffer's # structural detector flips the job to FAILED. # 2. ``adapter.error_events(exc)`` is the dispatcher's guaranteed - # fallback: emitted from the consumer side even when on_error itself - # raises (its body and the sentinel-put are wrapped in suppress so - # they cannot prevent shutdown but can fail silently). Without this - # yield, an on_error failure would leave the stream with no terminal - # error event at all and the buffer would mark the job COMPLETED. + # fallback: emitted from the consumer side when ``generate_flow_events`` + # raises before any cooperative terminal error reaches the queue. + # Without this yield, an early failure would leave the stream with no + # terminal error event and the buffer would mark the job COMPLETED. # 3. The buffer task's ``terminal_error_type`` check fires on either # RUN_ERROR source, so a single drive() failure cannot result in a # job marked COMPLETED. @@ -772,8 +771,6 @@ async def _stream_event_frames( raise except Exception as exc: # noqa: BLE001 drive_error = exc - with contextlib.suppress(Exception): - event_manager.on_error(data={"error": str(exc)}) with contextlib.suppress(Exception): await event_manager.queue.put((None, None, time.time())) # generate_flow_events emits on_end and the sentinel on success. @@ -789,6 +786,8 @@ async def _stream_event_frames( # AG-UI. A follow-up retires this once chat-view consumes AG-UI primitives. emit_side_channel = adapter.name == "agui" side_channel_events = frozenset({"add_message", "token", "remove_message", "error", "end"}) + terminal_error_type = getattr(adapter, "terminal_error_type", None) + terminal_error_seen = False seq = 0 run_task = asyncio.create_task(drive()) @@ -816,19 +815,22 @@ async def _stream_event_frames( ) seq += 1 for event in adapter.translate(event_type, event_data): + if terminal_error_type is not None and event.type == terminal_error_type: + terminal_error_seen = True yield _frame(event, seq) seq += 1 for event in adapter.final_events(): + if terminal_error_type is not None and event.type == terminal_error_type: + terminal_error_seen = True yield _frame(event, seq) seq += 1 # Guaranteed-fallback layer (see drive_error block above). If drive() - # captured an exception, emit the adapter's terminal error event(s) - # here even if on_error already produced one: an extra RUN_ERROR is - # cheap and only a hard error path will hit this. The duplicate is - # the cost of guaranteeing the consumer sees a terminal error event - # even when the cooperative on_error layer failed. - if drive_error is not None: + # captured an exception and no cooperative terminal error reached the + # stream, emit the adapter's terminal error event(s) here. + if drive_error is not None and not terminal_error_seen: for event in adapter.error_events(drive_error): + if terminal_error_type is not None and event.type == terminal_error_type: + terminal_error_seen = True yield _frame(event, seq) seq += 1 finally: @@ -872,6 +874,9 @@ def _execute_streaming_workflow( ) +_CancelEventsFactory = Callable[[str], Iterable[StreamEvent]] + + class _BackgroundRun: """In-memory buffer of a background run's protocol-native SSE frames for re-attach. @@ -889,8 +894,10 @@ class _BackgroundRun: from the new buffer head (replay loss is preferred over OOM). """ - def __init__(self, user_id: str) -> None: + def __init__(self, user_id: str, stream_protocol: str = "agui") -> None: self.user_id = user_id + self.stream_protocol = stream_protocol + self._cancel_events: _CancelEventsFactory | None = None self.frames: list[bytes] = [] # Index of the first frame still in ``frames`` (monotonic across the # life of the buffer). Once eviction starts, ``frames[i]`` corresponds @@ -899,13 +906,23 @@ class _BackgroundRun: self.done = False self._cond = asyncio.Condition() + def set_cancel_events(self, cancel_events: _CancelEventsFactory) -> None: + self._cancel_events = cancel_events + + @property + def has_cancel_events(self) -> bool: + return self._cancel_events is not None + + def _append_locked(self, frame: bytes) -> None: + self.frames.append(frame) + overflow = len(self.frames) - _MAX_FRAMES_PER_BACKGROUND_RUN + if overflow > 0: + del self.frames[:overflow] + self.base_index += overflow + async def append(self, frame: bytes) -> None: async with self._cond: - self.frames.append(frame) - overflow = len(self.frames) - _MAX_FRAMES_PER_BACKGROUND_RUN - if overflow > 0: - del self.frames[:overflow] - self.base_index += overflow + self._append_locked(frame) self._cond.notify_all() async def finish(self) -> None: @@ -913,6 +930,21 @@ class _BackgroundRun: self.done = True self._cond.notify_all() + async def finish_cancelled(self, reason: str) -> None: + """Append cancellation terminal events and finish replay exactly once.""" + async with self._cond: + if self.done: + return + events: list[StreamEvent] = [] + if self._cancel_events is not None: + with contextlib.suppress(Exception): + events = list(self._cancel_events(reason)) + for event in events: + seq = self.base_index + len(self.frames) + self._append_locked(format_sse_event(data_str=event.data_json, id=str(seq))) + self.done = True + self._cond.notify_all() + async def replay(self, start_index: int) -> AsyncIterator[bytes]: """Yield buffered frames from ``start_index`` and tail until done. @@ -953,6 +985,7 @@ _MAX_FRAMES_PER_BACKGROUND_RUN = 10_000 # letting frames accumulate without bound when the network is slow. _EVENT_QUEUE_MAX_SIZE = 256 _BACKGROUND_RUNS: dict[str, _BackgroundRun] = {} +_WORKFLOW_CANCELLED_MESSAGE = "Workflow run cancelled." async def _finalize_job_status(job_uuid: UUID, terminal_status: JobStatus) -> None: @@ -982,17 +1015,31 @@ async def _finalize_job_status(job_uuid: UUID, terminal_status: JobStatus) -> No async def _clear_background_run(job_id: str) -> None: """Pop the background run registry entry and wake any re-attach waiters. - Called from ``stop_workflow`` after revoking the buffer task. Without - this, ``_cond.wait()`` inside ``_BackgroundRun.replay`` can hang - indefinitely (the buffer is never marked done because the task that - would have called ``finish()`` was cancelled mid-execution), and the - cancelled ``_BackgroundRun`` lingers in memory until LRU evicts it. + Used for true cleanup paths that should discard replay state, such as + failed scheduling or test teardown. Cancellation normally finishes through + the owning buffer task so it can append protocol-native terminal events + before replay is marked done. """ bg_run = _BACKGROUND_RUNS.pop(job_id, None) if bg_run is not None: await bg_run.finish() +async def _finish_cancelled_background_run(job_id: str) -> None: + """Finish a stopped local run when no owner task is available to do it.""" + bg_run = _BACKGROUND_RUNS.get(job_id) + if bg_run is None or bg_run.done: + return + if not bg_run.has_cancel_events: + with contextlib.suppress(UnknownStreamProtocolError): + adapter = get_stream_adapter( + bg_run.stream_protocol, + StreamAdapterContext(run_id=job_id, thread_id=job_id), + ) + bg_run.set_cancel_events(adapter.cancel_events) + await bg_run.finish_cancelled(_WORKFLOW_CANCELLED_MESSAGE) + + def _register_background_run(job_id: str, bg_run: _BackgroundRun) -> None: """Register a background run, evicting a completed entry when full. @@ -1055,10 +1102,15 @@ async def _buffer_background_run( await _finalize_job_status(job_uuid, JobStatus.FAILED) return + bg_run.set_cancel_events(adapter.cancel_events) terminal_error_type = adapter.terminal_error_type fresh_background_tasks = BackgroundTasks() errored = False + cancelled = False + job_uuid = UUID(job_id) if isinstance(job_id, str) else job_id try: + with contextlib.suppress(Exception): + await get_job_service().update_job_status(job_uuid, JobStatus.IN_PROGRESS) async for frame, event_type in _stream_event_frames( adapter=adapter, flow_id=flow.id, @@ -1074,8 +1126,13 @@ async def _buffer_background_run( if terminal_error_type is not None and event_type == terminal_error_type: errored = True await bg_run.append(frame) + except asyncio.CancelledError: + cancelled = True + await bg_run.finish_cancelled(_WORKFLOW_CANCELLED_MESSAGE) + raise finally: - await bg_run.finish() + if not cancelled: + await bg_run.finish() # ``generate_flow_events`` queues telemetry, tracing teardown, and # other callbacks on this ``BackgroundTasks`` instance. In FastAPI's # request lifecycle those run after the response is sent; the @@ -1084,10 +1141,10 @@ async def _buffer_background_run( # not derail job-status finalization. with contextlib.suppress(Exception): await fresh_background_tasks() - # ``update_job_status`` queries the Job table by its UUID primary key; - # passing the raw string would silently miss every row. - job_uuid = UUID(job_id) if isinstance(job_id, str) else job_id - await _finalize_job_status(job_uuid, JobStatus.FAILED if errored else JobStatus.COMPLETED) + if not cancelled: + # ``update_job_status`` queries the Job table by its UUID primary key; + # passing the raw string would silently miss every row. + await _finalize_job_status(job_uuid, JobStatus.FAILED if errored else JobStatus.COMPLETED) async def execute_workflow_background( @@ -1117,7 +1174,7 @@ async def execute_workflow_background( user_id=current_user.id, ) - bg_run = _BackgroundRun(user_id=str(current_user.id)) + bg_run = _BackgroundRun(user_id=str(current_user.id), stream_protocol=stream_protocol) _register_background_run(job_id_str, bg_run) try: @@ -1408,9 +1465,10 @@ async def stop_workflow( }, ) await job_service.update_job_status(job_id, JobStatus.CANCELLED) - # Release the in-memory buffer and wake any re-attach waiters so they - # see a clean stream-end instead of hanging on ``_cond.wait()``. - await _clear_background_run(str(job_id)) + # The owning buffer task appends the protocol-native cancellation + # terminal event before marking replay done. This is only a fallback + # for races where a local buffer exists but no owner task is running. + await _finish_cancelled_background_run(str(job_id)) message = f"Job {job_id} cancelled successfully." return WorkflowStopResponse(job_id=str(job_id), message=message) diff --git a/src/backend/tests/unit/api/v2/adapters/test_agui_adapter.py b/src/backend/tests/unit/api/v2/adapters/test_agui_adapter.py index 4ba523f4af..83ce814199 100644 --- a/src/backend/tests/unit/api/v2/adapters/test_agui_adapter.py +++ b/src/backend/tests/unit/api/v2/adapters/test_agui_adapter.py @@ -120,6 +120,29 @@ class TestErrorEventsFallback: assert len(events) == 1 assert events[0].type == "RUN_ERROR" + def test_error_events_closes_open_text_message_before_run_error(self): + """Fallback errors must still leave the AG-UI text lifecycle well-formed.""" + adapter = get_stream_adapter("agui", _ctx()) + list(adapter.initial_events()) + list(adapter.translate("token", {"id": "m1", "chunk": "partial"})) + + events = list(adapter.error_events(RuntimeError("boom"))) + + assert [event.type for event in events] == ["TEXT_MESSAGE_END", "RUN_ERROR"] + assert json.loads(events[0].data_json)["messageId"] == "m1" + + def test_cancel_events_close_open_text_message_before_run_error(self): + """Cancellation uses adapter-owned framing and preserves AG-UI lifecycle.""" + adapter = get_stream_adapter("agui", _ctx()) + list(adapter.initial_events()) + list(adapter.translate("token", {"id": "m1", "chunk": "partial"})) + + events = list(adapter.cancel_events("cancelled")) + + assert [event.type for event in events] == ["TEXT_MESSAGE_END", "RUN_ERROR"] + assert json.loads(events[0].data_json)["messageId"] == "m1" + assert json.loads(events[1].data_json)["message"] == "cancelled" + class TestTerminalErrorType: """Buffer task watches for AG-UI RUN_ERROR to flip the job to FAILED.""" diff --git a/src/backend/tests/unit/api/v2/adapters/test_langflow_adapter.py b/src/backend/tests/unit/api/v2/adapters/test_langflow_adapter.py index a44ce267b1..3b32d1812c 100644 --- a/src/backend/tests/unit/api/v2/adapters/test_langflow_adapter.py +++ b/src/backend/tests/unit/api/v2/adapters/test_langflow_adapter.py @@ -62,6 +62,13 @@ class TestErrorHandling: assert payload == {"event": "error", "data": {"error": "boom"}} assert evt.type == "error" + def test_cancel_events_emit_protocol_error_payload(self): + adapter = get_stream_adapter("langflow", _ctx()) + [evt] = list(adapter.cancel_events("cancelled")) + payload = json.loads(evt.data_json) + assert payload == {"event": "error", "data": {"error": "cancelled"}} + assert evt.type == "error" + def test_terminal_error_type_is_error(self): """Used by the buffer task to decide JobStatus.FAILED.""" adapter = get_stream_adapter("langflow", _ctx()) diff --git a/src/backend/tests/unit/api/v2/adapters/test_registry.py b/src/backend/tests/unit/api/v2/adapters/test_registry.py index 6ed65e6024..04b2c22b3e 100644 --- a/src/backend/tests/unit/api/v2/adapters/test_registry.py +++ b/src/backend/tests/unit/api/v2/adapters/test_registry.py @@ -55,6 +55,9 @@ class _NoopAdapter: def error_events(self, _error): return [] + def cancel_events(self, _reason): + return [] + @property def terminal_error_type(self): return None diff --git a/src/backend/tests/unit/api/v2/test_agui_translator.py b/src/backend/tests/unit/api/v2/test_agui_translator.py index 0797105bc7..26b65ec041 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -90,20 +90,27 @@ def test_token_sequence_emits_start_contents_then_end_on_boundary(): assert isinstance(ended[1], RunFinishedEvent) -def test_new_message_id_closes_previous_message_and_opens_new(): +def test_interleaved_message_ids_remain_open_until_terminal_boundary(): t = AGUITranslator(run_id="r1", thread_id="t1") t.start() - t.translate("token", {"chunk": "a", "id": "m1"}) - out = t.translate("token", {"chunk": "b", "id": "m2"}) + first = t.translate("token", {"chunk": "a", "id": "m1"}) + second = t.translate("token", {"chunk": "b", "id": "m2"}) + resumed = t.translate("token", {"chunk": "c", "id": "m1"}) + ended = t.translate("end", {}) - assert isinstance(out[0], TextMessageEndEvent) - assert out[0].message_id == "m1" - assert isinstance(out[1], TextMessageStartEvent) - assert out[1].message_id == "m2" - assert isinstance(out[2], TextMessageContentEvent) - assert out[2].message_id == "m2" - assert out[2].delta == "b" + assert isinstance(first[0], TextMessageStartEvent) + assert first[0].message_id == "m1" + assert isinstance(second[0], TextMessageStartEvent) + assert second[0].message_id == "m2" + assert all(not isinstance(event, TextMessageEndEvent) for event in second) + assert len(resumed) == 1 + assert isinstance(resumed[0], TextMessageContentEvent) + assert resumed[0].message_id == "m1" + assert resumed[0].delta == "c" + ended_message_ids = {event.message_id for event in ended if isinstance(event, TextMessageEndEvent)} + assert ended_message_ids == {"m1", "m2"} + assert isinstance(ended[-1], RunFinishedEvent) def test_error_boundary_closes_open_text_message(): @@ -166,32 +173,49 @@ def test_token_for_already_ended_message_id_is_dropped(): ) -def test_token_after_boundary_close_is_dropped(): - """An interleaved token sequence ``A, B, A`` must not re-open id A. - - Switching from id A to id B closes A through ``_close_open_message``. A - later token for A used to slip past the dedup guard because - ``_close_open_message`` was the only finalizer that did not record the - closed id in ``_emitted_text_message_ids``. The translator must treat a - token-boundary close the same way it treats an ``add_message`` close. - """ +def test_token_after_interleaved_message_id_switch_is_preserved(): + """An interleaved token sequence ``A, B, A`` must preserve both streams.""" t = AGUITranslator(run_id="r1", thread_id="t1") t.start() a1 = t.translate("token", {"chunk": "hi", "id": "m1"}) assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m1" for e in a1) - # Switching to a new id closes m1 via _close_open_message. + # Switching to a new id opens it without closing m1; consumers group by id. b = t.translate("token", {"chunk": "yo", "id": "m2"}) - assert any(isinstance(e, TextMessageEndEvent) and e.message_id == "m1" for e in b) + assert all(not isinstance(e, TextMessageEndEvent) for e in b) assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m2" for e in b) - # A late token for m1 must be dropped, not re-open the ended message. + # A later token for m1 is still content for the already-open m1 lifecycle. late = t.translate("token", {"chunk": "again", "id": "m1"}) - assert all(not isinstance(e, TextMessageStartEvent) for e in late), ( - f"Token boundary did not mark m1 as ended; emitted: {late}" - ) - assert late == [], f"Expected no events for late token after boundary close; got {late}" + assert len(late) == 1 + assert isinstance(late[0], TextMessageContentEvent) + assert late[0].message_id == "m1" + assert late[0].delta == "again" + + +def test_partial_add_message_does_not_burn_streamed_text_id(): + """The first streaming chunk arrives as add_message before token events. + + ``add_message`` snapshots with state=partial must not emit a complete + TextMessage lifecycle, or later token events for the same id are dropped. + """ + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + partial = t.translate("add_message", {"id": "m1", "text": "Hel", "properties": {"state": "partial"}}) + first = t.translate("token", {"id": "m1", "chunk": "Hel"}) + second = t.translate("token", {"id": "m1", "chunk": "lo"}) + finalizer = t.translate("add_message", {"id": "m1", "text": "Hello", "properties": {"state": "complete"}}) + + text_lifecycle_types = (TextMessageStartEvent, TextMessageContentEvent, TextMessageEndEvent) + assert all(not isinstance(e, text_lifecycle_types) for e in partial) + assert [type(e) for e in first] == [TextMessageStartEvent, TextMessageContentEvent] + assert first[1].delta == "Hel" + assert len(second) == 1 + assert isinstance(second[0], TextMessageContentEvent) + assert second[0].delta == "lo" + assert [type(e) for e in finalizer] == [TextMessageEndEvent] def test_vertices_sorted_emits_state_snapshot_of_all_nodes(): diff --git a/src/backend/tests/unit/api/v2/test_workflow_agui.py b/src/backend/tests/unit/api/v2/test_workflow_agui.py index 7dd04af556..9633fd5065 100644 --- a/src/backend/tests/unit/api/v2/test_workflow_agui.py +++ b/src/backend/tests/unit/api/v2/test_workflow_agui.py @@ -44,6 +44,19 @@ def _agui_body(flow_id, *, message: str = "hello", mode: str = "sync", tweaks: d return body +def _sse_payloads(frames: list[bytes]) -> list[dict]: + return [ + json.loads(line.removeprefix("data:").strip()) + for frame in frames + for line in frame.decode("utf-8").splitlines() + if line.startswith("data:") + ] + + +def _sse_payload_type(payload: dict) -> str | None: + return payload.get("type") or payload.get("event") + + @pytest.fixture async def empty_flow(created_api_key): """Create a real empty flow owned by the API-key user; clean it up after.""" @@ -428,6 +441,84 @@ class TestAGUIStreaming: assert seen_maxsize == [2] assert event_types == ["token", "token", "error"] + @pytest.mark.parametrize(("stream_protocol", "terminal_type"), [("agui", "RUN_ERROR"), ("langflow", "error")]) + async def test_stream_error_path_emits_one_terminal_error( + self, + monkeypatch: pytest.MonkeyPatch, + stream_protocol: str, + terminal_type: str, + ): + """A producer that reports on_error then raises must not triple-emit terminal errors.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.adapters import StreamAdapterContext, get_stream_adapter + from langflow.api.v2.converters import ParsedWorkflowRun + + async def fake_generate_flow_events(**kwargs): + kwargs["event_manager"].on_error(data={"error": "inner"}) + message = "outer" + raise RuntimeError(message) + + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + adapter = get_stream_adapter( + stream_protocol, + StreamAdapterContext(run_id="run-1", thread_id="thread-1"), + ) + event_types = [ + event_type + async for _frame, event_type in workflow_module._stream_event_frames( + adapter=adapter, + flow_id=uuid4(), + flow_name="flow", + background_tasks=SimpleNamespace(add_task=lambda *_args, **_kwargs: None), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="", mode="stream"), + current_user=SimpleNamespace(id=uuid4()), + ) + ] + + assert event_types.count(terminal_type) == 1 + assert event_types[-1] == terminal_type + + @pytest.mark.parametrize(("stream_protocol", "terminal_type"), [("agui", "RUN_ERROR"), ("langflow", "error")]) + async def test_stream_error_path_uses_fallback_when_producer_raises_before_on_error( + self, + monkeypatch: pytest.MonkeyPatch, + stream_protocol: str, + terminal_type: str, + ): + """A producer that raises before on_error still emits one terminal error.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.adapters import StreamAdapterContext, get_stream_adapter + from langflow.api.v2.converters import ParsedWorkflowRun + + async def fake_generate_flow_events(**_kwargs): + message = "early boom" + raise RuntimeError(message) + + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + adapter = get_stream_adapter( + stream_protocol, + StreamAdapterContext(run_id="run-1", thread_id="thread-1"), + ) + frames = [ + frame + async for frame, _event_type in workflow_module._stream_event_frames( + adapter=adapter, + flow_id=uuid4(), + flow_name="flow", + background_tasks=SimpleNamespace(add_task=lambda *_args, **_kwargs: None), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="", mode="stream"), + current_user=SimpleNamespace(id=uuid4()), + ) + ] + + payloads = _sse_payloads(frames) + payload_types = [_sse_payload_type(payload) for payload in payloads] + assert payload_types.count(terminal_type) == 1 + assert payload_types[-1] == terminal_type + assert "early boom" in json.dumps(payloads[-1]) + async def test_agui_stream_emits_end_side_channel_for_build_duration(self, monkeypatch: pytest.MonkeyPatch): """The AG-UI stream must preserve v1 end payloads for chat build-duration persistence.""" from langflow.api.v2 import workflow as workflow_module @@ -885,6 +976,9 @@ class TestAGUIBackgroundJobStatus: class FakeAdapter: terminal_error_type = "RUN_ERROR" + def cancel_events(self, _reason): + return [] + async def fake_stream_event_frames(**kwargs): captured.update(kwargs) yield b"data: {}\n\n", "RUN_FINISHED" @@ -906,6 +1000,234 @@ class TestAGUIBackgroundJobStatus: assert captured["run_id"] == str(job_id) assert captured["track_job_status"] is False + async def test_background_buffer_marks_job_in_progress(self, monkeypatch: pytest.MonkeyPatch): + """A background job should leave QUEUED once its buffer task starts executing.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + job_id = uuid4() + updates: list[tuple[object, object, bool]] = [] + + class FakeAdapter: + name = "agui" + terminal_error_type = "RUN_ERROR" + + def cancel_events(self, _reason): + return [] + + class FakeJobService: + async def update_job_status(self, seen_job_id, status, *, finished_timestamp=False): + updates.append((seen_job_id, status, finished_timestamp)) + + async def fake_stream_event_frames(**_kwargs): + yield b"data: {}\n\n", "RUN_FINISHED" + + monkeypatch.setattr(workflow_module, "get_stream_adapter", lambda *_args, **_kwargs: FakeAdapter()) + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + monkeypatch.setattr(workflow_module, "_stream_event_frames", fake_stream_event_frames) + monkeypatch.setattr(workflow_module, "_finalize_job_status", AsyncMock()) + + bg_run = workflow_module._BackgroundRun(user_id=str(uuid4())) + await workflow_module._buffer_background_run( + bg_run=bg_run, + flow=SimpleNamespace(id=uuid4(), name="flow"), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="hi", mode="background"), + job_id=str(job_id), + current_user=SimpleNamespace(id=uuid4()), + stream_protocol="agui", + ) + + assert (job_id, workflow_module.JobStatus.IN_PROGRESS, False) in updates + + async def test_cancelled_agui_buffer_wakes_tail_reader_with_closed_text_and_run_error( + self, monkeypatch: pytest.MonkeyPatch + ): + """The owner task must append cancellation before marking replay done.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + started = asyncio.Event() + + class FakeJobService: + async def update_job_status(self, *_args, **_kwargs): + return None + + async def fake_generate_flow_events(**kwargs): + kwargs["event_manager"].on_token(data={"id": "m1", "chunk": "partial"}) + started.set() + await asyncio.Event().wait() + + async def collect_tail(bg_run): + tail = bg_run.base_index + len(bg_run.frames) + return [frame async for frame in bg_run.replay(tail)] + + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + monkeypatch.setattr(workflow_module, "_finalize_job_status", AsyncMock()) + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + bg_run = workflow_module._BackgroundRun(user_id=str(uuid4()), stream_protocol="agui") + buffer_task = asyncio.create_task( + workflow_module._buffer_background_run( + bg_run=bg_run, + flow=SimpleNamespace(id=uuid4(), name="flow"), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="hi", mode="background"), + job_id=str(uuid4()), + current_user=SimpleNamespace(id=uuid4()), + stream_protocol="agui", + ) + ) + await asyncio.wait_for(started.wait(), timeout=2) + for _ in range(100): + payload_types = [_sse_payload_type(payload) for payload in _sse_payloads(bg_run.frames)] + if "TEXT_MESSAGE_CONTENT" in payload_types: + break + await asyncio.sleep(0.01) + else: + pytest.fail("token content was never buffered") + + tail_task = asyncio.create_task(collect_tail(bg_run)) + await asyncio.sleep(0) + + buffer_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await buffer_task + tail_frames = await asyncio.wait_for(tail_task, timeout=2) + + tail_payloads = _sse_payloads(tail_frames) + assert [_sse_payload_type(payload) for payload in tail_payloads] == ["TEXT_MESSAGE_END", "RUN_ERROR"] + assert tail_payloads[0]["messageId"] == "m1" + assert tail_payloads[1]["message"] == "Workflow run cancelled." + + async def test_cancelled_langflow_buffer_wakes_tail_reader_with_langflow_error( + self, monkeypatch: pytest.MonkeyPatch + ): + """Cancellation framing must stay protocol-native outside AG-UI too.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + started = asyncio.Event() + + class FakeJobService: + async def update_job_status(self, *_args, **_kwargs): + return None + + async def fake_generate_flow_events(**kwargs): + kwargs["event_manager"].on_token(data={"id": "m1", "chunk": "partial"}) + started.set() + await asyncio.Event().wait() + + async def collect_tail(bg_run): + tail = bg_run.base_index + len(bg_run.frames) + return [frame async for frame in bg_run.replay(tail)] + + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + monkeypatch.setattr(workflow_module, "_finalize_job_status", AsyncMock()) + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + + bg_run = workflow_module._BackgroundRun(user_id=str(uuid4()), stream_protocol="langflow") + buffer_task = asyncio.create_task( + workflow_module._buffer_background_run( + bg_run=bg_run, + flow=SimpleNamespace(id=uuid4(), name="flow"), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="hi", mode="background"), + job_id=str(uuid4()), + current_user=SimpleNamespace(id=uuid4()), + stream_protocol="langflow", + ) + ) + await asyncio.wait_for(started.wait(), timeout=2) + for _ in range(100): + payload_types = [_sse_payload_type(payload) for payload in _sse_payloads(bg_run.frames)] + if "token" in payload_types: + break + await asyncio.sleep(0.01) + else: + pytest.fail("token event was never buffered") + + tail_task = asyncio.create_task(collect_tail(bg_run)) + await asyncio.sleep(0) + + buffer_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await buffer_task + tail_frames = await asyncio.wait_for(tail_task, timeout=2) + + [payload] = _sse_payloads(tail_frames) + assert payload == {"event": "error", "data": {"error": "Workflow run cancelled."}} + + async def test_finish_cancelled_background_run_appends_terminal_before_waking_replay( + self, monkeypatch: pytest.MonkeyPatch + ): + """The stop fallback must not wake replay readers before cancellation is buffered.""" + from langflow.api.v2 import workflow as workflow_module + from langflow.api.v2.converters import ParsedWorkflowRun + + job_id = str(uuid4()) + started = asyncio.Event() + + class FakeJobService: + async def update_job_status(self, *_args, **_kwargs): + return None + + async def fake_generate_flow_events(**kwargs): + kwargs["event_manager"].on_token(data={"id": "m1", "chunk": "partial"}) + started.set() + await asyncio.Event().wait() + + async def collect_tail(bg_run): + tail = bg_run.base_index + len(bg_run.frames) + return [frame async for frame in bg_run.replay(tail)] + + monkeypatch.setattr(workflow_module, "get_job_service", lambda: FakeJobService()) + monkeypatch.setattr(workflow_module, "_finalize_job_status", AsyncMock()) + monkeypatch.setattr(workflow_module, "generate_flow_events", fake_generate_flow_events) + monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {}) + + bg_run = workflow_module._BackgroundRun(user_id=str(uuid4()), stream_protocol="agui") + workflow_module._BACKGROUND_RUNS[job_id] = bg_run + buffer_task = asyncio.create_task( + workflow_module._buffer_background_run( + bg_run=bg_run, + flow=SimpleNamespace(id=uuid4(), name="flow"), + parsed=ParsedWorkflowRun(flow_id=str(uuid4()), input_value="hi", mode="background"), + job_id=job_id, + current_user=SimpleNamespace(id=uuid4()), + stream_protocol="agui", + ) + ) + try: + await asyncio.wait_for(started.wait(), timeout=2) + for _ in range(100): + payload_types = [_sse_payload_type(payload) for payload in _sse_payloads(bg_run.frames)] + if "TEXT_MESSAGE_CONTENT" in payload_types: + break + await asyncio.sleep(0.01) + else: + pytest.fail("token content was never buffered") + + tail_task = asyncio.create_task(collect_tail(bg_run)) + await asyncio.sleep(0) + + await workflow_module._finish_cancelled_background_run(job_id) + tail_frames = await asyncio.wait_for(tail_task, timeout=2) + + tail_payloads = _sse_payloads(tail_frames) + assert [_sse_payload_type(payload) for payload in tail_payloads] == ["TEXT_MESSAGE_END", "RUN_ERROR"] + assert tail_payloads[0]["messageId"] == "m1" + assert tail_payloads[1]["message"] == "Workflow run cancelled." + + frames_after_fallback = list(bg_run.frames) + buffer_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await buffer_task + assert bg_run.frames == frames_after_fallback + finally: + if not buffer_task.done(): + buffer_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await buffer_task + await workflow_module._clear_background_run(job_id) + async def test_message_text_containing_run_error_literal_does_not_fail_job( self, client: AsyncClient, @@ -1956,18 +2278,19 @@ class TestExecuteWorkflowBackgroundQueueOwnership: assert created_jobs == [job_id] assert started_jobs == [str(job_id)] assert registered_owners == [] + assert response.links["events"] == f"/api/v2/workflows/{job_id}/events" class TestStopWorkflowEndToEnd: - """The full ``POST /workflows/stop`` HTTP flow clears the in-memory buffer too. + """The full ``POST /workflows/stop`` HTTP flow terminates the in-memory buffer too. ``test_background_run_can_be_stopped`` covers the 200 response. This class - pins the side-effects: the in-process ``_BACKGROUND_RUNS`` entry is popped, - the ``_BackgroundRun.done`` flag flips so re-attach readers unblock, and - the Job row is marked ``CANCELLED``. + pins the side-effects: the in-process ``_BACKGROUND_RUNS`` entry remains + replayable, the ``_BackgroundRun.done`` flag flips so re-attach readers + unblock, and the Job row is marked ``CANCELLED``. """ - async def test_stop_clears_in_memory_registry_and_marks_job_cancelled( + async def test_stop_finishes_replay_buffer_and_marks_job_cancelled( self, client: AsyncClient, created_api_key, @@ -1987,27 +2310,37 @@ class TestStopWorkflowEndToEnd: assert start.status_code == 200 job_id = start.json()["job_id"] - # The registry holds the buffer keyed by job_id until /stop or finish. + # The registry holds the buffer keyed by job_id so reconnect can replay + # the cancellation terminal event after /stop returns. assert job_id in workflow_module._BACKGROUND_RUNS, ( "Background run was not registered; the stop assertions below would be vacuous" ) bg_run = workflow_module._BACKGROUND_RUNS[job_id] - stop = await client.post( - "api/v2/workflows/stop", - json={"job_id": job_id}, - headers=headers, - ) - assert stop.status_code == 200 + try: + stop = await client.post( + "api/v2/workflows/stop", + json={"job_id": job_id}, + headers=headers, + ) + assert stop.status_code == 200 - # Side-effects: registry popped, buffer finished, job row CANCELLED. - assert job_id not in workflow_module._BACKGROUND_RUNS - assert bg_run.done is True + # Side-effects: registry remains replayable, buffer finished, job row CANCELLED. + assert job_id in workflow_module._BACKGROUND_RUNS + assert bg_run.done is True - async with session_scope() as session: - row = await session.get(Job, _UUID(job_id)) - assert row is not None - assert row.status == JobStatus.CANCELLED + events = await client.get(f"api/v2/workflows/{job_id}/events", headers=headers) + assert events.status_code == 200 + assert "RUN_ERROR" in events.text + assert "cancelled" in events.text.lower() + assert "RUN_FINISHED" not in events.text + + async with session_scope() as session: + row = await session.get(Job, _UUID(job_id)) + assert row is not None + assert row.status == JobStatus.CANCELLED + finally: + await workflow_module._clear_background_run(job_id) async def test_stop_cancels_queue_owned_background_task_when_task_service_noops( self, diff --git a/src/lfx/src/lfx/custom/custom_component/component.py b/src/lfx/src/lfx/custom/custom_component/component.py index d9cd99a027..5c246e7cc6 100644 --- a/src/lfx/src/lfx/custom/custom_component/component.py +++ b/src/lfx/src/lfx/custom/custom_component/component.py @@ -2061,8 +2061,9 @@ class Component(CustomComponent): if self._event_manager: if first_chunk: # Send the initial message only on the first chunk - msg_copy = message.model_copy() + msg_copy = message.model_copy(update={"properties": message.properties.model_copy(deep=True)}) msg_copy.text = complete_message + msg_copy.properties.state = "partial" await self._send_message_event(msg_copy, id_=message_id) await asyncio.to_thread( self._event_manager.on_token, diff --git a/src/lfx/src/lfx/schema/workflow.py b/src/lfx/src/lfx/schema/workflow.py index 0e19a243aa..11dbe91f03 100644 --- a/src/lfx/src/lfx/schema/workflow.py +++ b/src/lfx/src/lfx/schema/workflow.py @@ -423,6 +423,7 @@ class WorkflowJobResponse(BaseModel): if not self.links: self.links = { "status": f"/api/v2/workflows?job_id={self.job_id!s}", + "events": f"/api/v2/workflows/{self.job_id!s}/events", "stop": "/api/v2/workflows/stop", } return self diff --git a/src/lfx/tests/unit/custom/custom_component/test_component.py b/src/lfx/tests/unit/custom/custom_component/test_component.py index b7058f96db..719fd1ecb5 100644 --- a/src/lfx/tests/unit/custom/custom_component/test_component.py +++ b/src/lfx/tests/unit/custom/custom_component/test_component.py @@ -140,6 +140,31 @@ async def test_send_message_without_database(): assert event_manager.on_message.called +@pytest.mark.asyncio +async def test_process_chunk_marks_first_event_partial_without_mutating_message(): + component = Component() + messages: list[dict] = [] + tokens: list[dict] = [] + + class EventManager: + def on_message(self, data): + messages.append(data) + + def on_token(self, data): + tokens.append(data) + + component._event_manager = EventManager() + message = Message(text="", sender="AI") + message.properties.state = "complete" + + complete_message = await component._process_chunk("", "", "message-id", message, first_chunk=True) + + assert complete_message == "" + assert message.properties.state == "complete" + assert messages[0]["properties"]["state"] == "partial" + assert tokens == [{"chunk": "", "id": "message-id"}] + + @pytest.mark.usefixtures("use_noop_session") @pytest.mark.asyncio async def test_agent_component_send_message_events(monkeypatch): # noqa: ARG001 From 3a124dd711ed7582a187808525f6bcc59eba9ef6 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 11 Jun 2026 14:58:59 +0000 Subject: [PATCH 08/11] [autofix.ci] apply automated fixes --- .../API/agui/__tests__/apply-state-delta.test.ts | 6 +----- .../agui/__tests__/run-flow-bridge-e2e.test.ts | 4 +--- .../src/controllers/API/agui/run-flow-bridge.ts | 15 ++++++++++----- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts index b6a7bd1427..d92480b011 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts @@ -319,11 +319,7 @@ describe("applyStateDelta", () => { useFlowStore.setState({ clearAndSetEdgesRunning: original }); - expect(calls).toEqual([ - ["node-a"], - ["node-a", "node-b"], - ["node-b"], - ]); + expect(calls).toEqual([["node-a"], ["node-a", "node-b"], ["node-b"]]); }); }); }); diff --git a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts index 211f68ff3c..d46f62a58e 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/run-flow-bridge-e2e.test.ts @@ -612,9 +612,7 @@ describe("runFlowAGUI end-to-end", () => { expect(signals[0].aborted).toBe(true); const flow = useFlowStore.getState(); expect(flow.isBuilding).toBe(false); - expect(flow.flowBuildStatus["node-a"]?.status).toBe( - BuildStatus.TO_BUILD, - ); + expect(flow.flowBuildStatus["node-a"]?.status).toBe(BuildStatus.TO_BUILD); expect(flow.flowBuildStatus["node-b"]?.status).toBe(BuildStatus.BUILT); // ``buildInfo`` is recorded as a failure (so trackFlowBuild sees the // cancellation as an error) but without an ``error`` string: the diff --git a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts index 5dc576c6c8..000967a206 100644 --- a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts +++ b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts @@ -297,11 +297,16 @@ export async function runFlowAGUI( }; const appendLogEvent = (data: unknown) => { - const { component_id, output, name, message, type: logType } = - (data ?? {}) as { - component_id?: string; - output?: string; - } & Partial; + const { + component_id, + output, + name, + message, + type: logType, + } = (data ?? {}) as { + component_id?: string; + output?: string; + } & Partial; if (!component_id || !output) { console.error( "[runFlowAGUI] Received malformed log event; missing component_id or output", From 78baed8104d163008dfe09d805e961f4dd9b8d5a Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Mon, 15 Jun 2026 12:04:19 -0300 Subject: [PATCH 09/11] fix(frontend): enable downlevelIteration for jest Set/Map iteration ts-jest compiles with target es5; without downlevelIteration, [...set] and for...of over a Set/Map emit ES5 that yields nothing. That silently broke the AG-UI bridge tests: runningNodeIds spread, markRunningNodesFailed, and restoreOriginalBuildStatuses all iterated empty. Production (Vite/SWC, modern target) was never affected; only the ts-jest harness was. Fixes the 3 failing jest tests on this branch with no other suite changes (4994/4994 pass). --- src/frontend/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/src/frontend/tsconfig.json b/src/frontend/tsconfig.json index d7adb91480..cea6e03790 100644 --- a/src/frontend/tsconfig.json +++ b/src/frontend/tsconfig.json @@ -1,5 +1,6 @@ { "compilerOptions": { + "downlevelIteration": true, "target": "es5", "lib": ["dom", "dom.iterable", "esnext"], "allowJs": true, From 76bc8ee03560135b79c9800dc0dd827c11f4445a Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Mon, 15 Jun 2026 12:12:21 -0300 Subject: [PATCH 10/11] fix(api/v2): surface inactivated branch vertices over AG-UI A branch component (If-Else, Conditional Router) reports its not-taken vertices in build_data.inactivated_vertices, but the AG-UI translator only emitted the branch node's own success/error status and dropped that list. The canvas seeds every planned node as pending from vertices_sorted; skipped vertices then get no build_start/end_vertex, so they stayed stuck on pending instead of rendering as inactive (the v1 build path marked them INACTIVE). The translator now appends an inactive STATE_DELTA op per inactivated vertex, and the frontend bridge maps the new inactive status to BuildStatus.INACTIVE and tears its edges down like a completed node. Fixes the If-Else regression in general-bugs-reset-flow-run.spec.ts. --- .../base/langflow/api/v2/agui_translator.py | 13 +++++++- .../tests/unit/api/v2/test_agui_translator.py | 33 +++++++++++++++++++ .../agui/__tests__/apply-state-delta.test.ts | 21 ++++++++++++ .../controllers/API/agui/run-flow-bridge.ts | 9 ++++- 4 files changed, 74 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/api/v2/agui_translator.py b/src/backend/base/langflow/api/v2/agui_translator.py index 82f7ba895c..d6672da2bb 100644 --- a/src/backend/base/langflow/api/v2/agui_translator.py +++ b/src/backend/base/langflow/api/v2/agui_translator.py @@ -209,9 +209,20 @@ class AGUITranslator: if not node_id: return [] status = "success" if build_data.get("valid") else "error" + delta = [self._set_node(node_id, status, build_data.get("data"))] + # A branch component (If-Else, Conditional Router) reports the vertices on + # the not-taken branch in ``inactivated_vertices`` (already unioned with the + # transitively excluded set in ``build.py``). The canvas seeds those nodes as + # ``pending`` from ``vertices_sorted`` and never builds them, so without this + # they stay stuck on ``pending`` instead of rendering as inactive. Mark each + # one ``inactive`` so the frontend maps it to ``BuildStatus.INACTIVE``. + delta.extend( + self._set_node(inactive_id, "inactive", None) + for inactive_id in build_data.get("inactivated_vertices") or [] + ) return [ StepFinishedEvent(step_name=node_id), - StateDeltaEvent(delta=[self._set_node(node_id, status, build_data.get("data"))]), + StateDeltaEvent(delta=delta), ] def _translate_add_message(self, data: dict) -> list[BaseEvent]: diff --git a/src/backend/tests/unit/api/v2/test_agui_translator.py b/src/backend/tests/unit/api/v2/test_agui_translator.py index 890e3f6028..0a29d94a20 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -487,6 +487,39 @@ def test_end_vertex_failure_sets_error_status(): assert op["value"]["status"] == "error" +def test_end_vertex_marks_inactivated_vertices_inactive(): + """A branch component's not-taken vertices must be marked ``inactive``. + + Without this the canvas leaves them on the ``pending`` status seeded by + ``vertices_sorted`` (they are skipped, so no build_start/end_vertex follows), + and the frontend never renders the inactive/skipped state for conditional + routing (If-Else, Conditional Router). + """ + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + out = t.translate( + "end_vertex", + { + "build_data": { + "id": "branch", + "valid": True, + "data": None, + "inactivated_vertices": ["node-a", "node-b"], + } + }, + ) + + ops = out[1].delta + assert ops[0]["path"] == "/nodes/branch" + assert ops[0]["value"]["status"] == "success" + inactive = {op["path"]: op["value"]["status"] for op in ops[1:]} + assert inactive == { + "/nodes/node-a": "inactive", + "/nodes/node-b": "inactive", + } + + def test_node_state_deltas_use_add_so_they_apply_without_vertices_sorted(): """build_start/end_vertex must not depend on vertices_sorted seeding the node. diff --git a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts index d92480b011..16a3739585 100644 --- a/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts +++ b/src/frontend/src/controllers/API/agui/__tests__/apply-state-delta.test.ts @@ -165,6 +165,27 @@ describe("applyStateDelta", () => { expect(touched.has("node-a")).toBe(true); }); + it("maps inactive status to BuildStatus.INACTIVE and skips the flow pool", () => { + const touched = new Set(); + + applyStateDelta( + [ + { + op: "add", + path: "/nodes/node-a", + value: { status: "inactive", output: null }, + }, + ], + "run-1", + touched, + ); + + const state = useFlowStore.getState(); + expect(state.flowBuildStatus["node-a"]?.status).toBe(BuildStatus.INACTIVE); + expect(state.flowPool["node-a"]).toBeUndefined(); + expect(touched.has("node-a")).toBe(true); + }); + it("falls back to BuildStatus.BUILDING when the status is unknown", () => { const touched = new Set(); diff --git a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts index 000967a206..c4af8dde24 100644 --- a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts +++ b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts @@ -38,6 +38,9 @@ const AGUI_STATUS_TO_BUILD_STATUS: Record = { running: BuildStatus.BUILDING, success: BuildStatus.BUILT, error: BuildStatus.ERROR, + // A branch component's not-taken vertices arrive as ``inactive`` so the + // canvas can render them as skipped instead of leaving them on ``pending``. + inactive: BuildStatus.INACTIVE, }; interface JsonPatchOp { @@ -168,7 +171,11 @@ export function applyStateDelta( } else { flowStore.updateEdgesRunningByNodes([nodeId], true); } - } else if (value.status === "success" || value.status === "error") { + } else if ( + value.status === "success" || + value.status === "error" || + value.status === "inactive" + ) { if (runningNodeIds) { runningNodeIds.delete(nodeId); flowStore.clearAndSetEdgesRunning([...runningNodeIds]); From 91f650c7639e9e18952d68f883abd3be2bcbbb7d Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Mon, 15 Jun 2026 14:11:27 -0300 Subject: [PATCH 11/11] fix(api/v2): dedupe repeated inactive node deltas in AG-UI stream build.py keeps reporting a conditionally-excluded vertex in inactivated_vertices on every subsequent end_vertex (the excluded set persists until the ConditionalRouter clears it), so the translator was putting the same inactive STATE_DELTA on the wire once per remaining vertex. Track emitted inactive nodes and skip re-emitting; drop a node from the set when it actually runs again (build_start/end_vertex) so a loop re-activation can still re-emit inactive later. --- .../base/langflow/api/v2/agui_translator.py | 25 ++++++++-- .../tests/unit/api/v2/test_agui_translator.py | 46 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/src/backend/base/langflow/api/v2/agui_translator.py b/src/backend/base/langflow/api/v2/agui_translator.py index d6672da2bb..60c447e07d 100644 --- a/src/backend/base/langflow/api/v2/agui_translator.py +++ b/src/backend/base/langflow/api/v2/agui_translator.py @@ -88,6 +88,14 @@ class AGUITranslator: # Custom content blocks already emitted as CUSTOM events, mapped to a # fingerprint of their last-emitted payload so in-place updates re-emit. self._emitted_content_state: dict[str, str] = {} + # Nodes already emitted as ``inactive``. ``build.py`` keeps reporting a + # conditionally-excluded vertex in ``inactivated_vertices`` on every + # subsequent ``end_vertex`` (the excluded set persists until the router + # clears it), so without this the same inactive delta goes out once per + # remaining vertex. A node is dropped from this set when it actually runs + # again (``build_start``/``end_vertex`` for it), so a re-activation in a + # loop can re-emit ``inactive`` later. + self._inactivated_nodes: set[str] = set() def start(self) -> list[BaseEvent]: """Open the run. @@ -197,6 +205,8 @@ class AGUITranslator: node_id = data.get("id") if not node_id: return [] + # The node is running again, so a later exclusion may re-emit ``inactive``. + self._inactivated_nodes.discard(node_id) return [ StepStartedEvent(step_name=node_id), StateDeltaEvent(delta=[self._set_node(node_id, "running", None)]), @@ -209,17 +219,22 @@ class AGUITranslator: if not node_id: return [] status = "success" if build_data.get("valid") else "error" + # This node just ran, so a later exclusion may re-emit ``inactive``. + self._inactivated_nodes.discard(node_id) delta = [self._set_node(node_id, status, build_data.get("data"))] # A branch component (If-Else, Conditional Router) reports the vertices on # the not-taken branch in ``inactivated_vertices`` (already unioned with the # transitively excluded set in ``build.py``). The canvas seeds those nodes as # ``pending`` from ``vertices_sorted`` and never builds them, so without this # they stay stuck on ``pending`` instead of rendering as inactive. Mark each - # one ``inactive`` so the frontend maps it to ``BuildStatus.INACTIVE``. - delta.extend( - self._set_node(inactive_id, "inactive", None) - for inactive_id in build_data.get("inactivated_vertices") or [] - ) + # one ``inactive`` so the frontend maps it to ``BuildStatus.INACTIVE``, but + # only once: ``build.py`` re-reports the persisted excluded set on every + # later ``end_vertex``, so dedupe to avoid sending the same delta repeatedly. + for inactive_id in build_data.get("inactivated_vertices") or []: + if inactive_id in self._inactivated_nodes: + continue + self._inactivated_nodes.add(inactive_id) + delta.append(self._set_node(inactive_id, "inactive", None)) return [ StepFinishedEvent(step_name=node_id), StateDeltaEvent(delta=delta), diff --git a/src/backend/tests/unit/api/v2/test_agui_translator.py b/src/backend/tests/unit/api/v2/test_agui_translator.py index 0a29d94a20..6d237e8cfc 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -520,6 +520,52 @@ def test_end_vertex_marks_inactivated_vertices_inactive(): } +def test_inactivated_vertex_is_emitted_once_across_end_vertices(): + """``build.py`` keeps reporting the persisted excluded set on later end_vertex. + + The translator must emit each node's ``inactive`` delta only once so the first + run of the feature doesn't put the same redundant op on the wire repeatedly. + """ + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + first = t.translate( + "end_vertex", + {"build_data": {"id": "router", "valid": True, "data": None, "inactivated_vertices": ["node-b"]}}, + ) + # The next vertex still carries node-b in the persisted excluded set. + second = t.translate( + "end_vertex", + {"build_data": {"id": "node-a", "valid": True, "data": None, "inactivated_vertices": ["node-b"]}}, + ) + + def inactive_paths(out): + return [op["path"] for op in out[1].delta if op["value"]["status"] == "inactive"] + + assert inactive_paths(first) == ["/nodes/node-b"] + assert inactive_paths(second) == [] # already emitted, not repeated + + +def test_reactivated_vertex_can_emit_inactive_again(): + """A vertex that runs after being excluded (loop re-activation) may go inactive again.""" + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + excluded = t.translate( + "end_vertex", + {"build_data": {"id": "router", "valid": True, "data": None, "inactivated_vertices": ["node-b"]}}, + ) + assert [op["path"] for op in excluded[1].delta if op["value"]["status"] == "inactive"] == ["/nodes/node-b"] + + # node-b actually runs on a later pass, then gets excluded again. + t.translate("build_start", {"id": "node-b"}) + again = t.translate( + "end_vertex", + {"build_data": {"id": "router", "valid": True, "data": None, "inactivated_vertices": ["node-b"]}}, + ) + assert [op["path"] for op in again[1].delta if op["value"]["status"] == "inactive"] == ["/nodes/node-b"] + + def test_node_state_deltas_use_add_so_they_apply_without_vertices_sorted(): """build_start/end_vertex must not depend on vertices_sorted seeding the node.