diff --git a/src/backend/base/langflow/api/build.py b/src/backend/base/langflow/api/build.py index b703e72c9b..37b0b1ec86 100644 --- a/src/backend/base/langflow/api/build.py +++ b/src/backend/base/langflow/api/build.py @@ -363,6 +363,8 @@ 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, + track_job_status: bool = True, ) -> None: """Generate events for flow building process. @@ -370,6 +372,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 +386,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 +406,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 +529,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 +543,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) @@ -676,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/adapters/__init__.py b/src/backend/base/langflow/api/v2/adapters/__init__.py index 1f2dc879b9..9a670666bc 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 ffe63379f0..742faee3d2 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, @@ -77,7 +77,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 4a0417eb04..753e85e914 100644 --- a/src/backend/base/langflow/api/v2/adapters/langflow.py +++ b/src/backend/base/langflow/api/v2/adapters/langflow.py @@ -104,6 +104,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 b8db8e7120..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,9 +219,25 @@ 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``, 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=[self._set_node(node_id, status, build_data.get("data"))]), + StateDeltaEvent(delta=delta), ] def _translate_add_message(self, data: dict) -> list[BaseEvent]: @@ -263,7 +289,7 @@ class AGUITranslator: # 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 not is_partial and message_id not in self._emitted_text_message_ids: if self._open_message_id is None: self._emitted_text_message_ids.add(message_id) events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) 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/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 6e9e2e5f72..6d237e8cfc 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -284,6 +284,26 @@ def test_error_drains_buffered_messages_before_run_error(): assert deltas["m3"] == "second waiting" +def test_partial_add_message_before_token_does_not_burn_streamed_text_id(): + """A partial snapshot before token streaming must not complete the id.""" + 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_partial_add_message_does_not_finalize_open_message(): """A partial ``add_message`` re-fire must not close the streaming message. @@ -467,6 +487,85 @@ 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_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. diff --git a/src/backend/tests/unit/api/v2/test_workflow.py b/src/backend/tests/unit/api/v2/test_workflow.py index e4cba4dbf3..b8d38bb1dd 100644 --- a/src/backend/tests/unit/api/v2/test_workflow.py +++ b/src/backend/tests/unit/api/v2/test_workflow.py @@ -287,16 +287,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) @@ -305,9 +303,48 @@ 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_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, @@ -580,15 +617,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_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): 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..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(); @@ -218,7 +239,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 +303,44 @@ 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..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 @@ -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,8 @@ 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..c4af8dde24 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"; @@ -31,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 { @@ -44,6 +54,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 +68,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 +99,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 +135,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 +149,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 +165,23 @@ export function applyStateDelta( // bridge's final ``finish()`` clears edges, and they never animate during // the run. if (value.status === "running") { - flowStore.updateEdgesRunningByNodes([nodeId], true); - } else if (value.status === "success" || value.status === "error") { - flowStore.updateEdgesRunningByNodes([nodeId], false); + if (runningNodeIds) { + runningNodeIds.add(nodeId); + flowStore.clearAndSetEdgesRunning([...runningNodeIds]); + } else { + flowStore.updateEdgesRunningByNodes([nodeId], true); + } + } else if ( + value.status === "success" || + value.status === "error" || + value.status === "inactive" + ) { + 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 +249,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 +265,96 @@ 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 +378,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 +408,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 +458,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 +470,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, { 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, 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/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