682 client - Add support for asynchronous workflow test execution and SSE streaming

Author: ivicac

client/src/ee/pages/embedded/automation-workflows/workflow-builder/components/workflow-builder-header/hooks/useWorkflowBuilderHeader.ts

@@ -7,10 +7,12 @@ import useWorkflowEditorStore from '@/pages/platform/workflow-editor/stores/useW
 import useWorkflowNodeDetailsPanelStore from '@/pages/platform/workflow-editor/stores/useWorkflowNodeDetailsPanelStore';
 import useWorkflowTestChatStore from '@/pages/platform/workflow-editor/stores/useWorkflowTestChatStore';
 import {useAnalytics} from '@/shared/hooks/useAnalytics';
+import {useWorkflowTestStream} from '@/shared/hooks/useWorkflowTestStream';
 import {WorkflowTestApi} from '@/shared/middleware/platform/workflow/test';
 import {useEnvironmentStore} from '@/shared/stores/useEnvironmentStore';
+import {getTestWorkflowAttachRequest, getTestWorkflowStreamPostRequest} from '@/shared/util/testWorkflow-utils';
 import {useQueryClient} from '@tanstack/react-query';
-import {RefObject, useCallback, useEffect} from 'react';
+import {RefObject, useCallback, useEffect, useState} from 'react';
 import {ImperativePanelHandle} from 'react-resizable-panels';
 import {useNavigate, useParams, useSearchParams} from 'react-router-dom';
 import {useShallow} from 'zustand/react/shallow';
@@ -24,6 +26,7 @@ interface UseProjectHeaderProps {
 }
 
 export const useWorkflowBuilderHeader = ({bottomResizablePanelRef, chatTrigger, projectId}: UseProjectHeaderProps) => {
+    const [jobId, setJobId] = useState<string | null>(null);
     const setDataPillPanelOpen = useDataPillPanelStore((state) => state.setDataPillPanelOpen);
     const currentEnvironmentId = useEnvironmentStore((state) => state.currentEnvironmentId);
     const workflow = useWorkflowDataStore((state) => state.workflow);
@@ -98,7 +101,20 @@ export const useWorkflowBuilderHeader = ({bottomResizablePanelRef, chatTrigger,
         );
     };
 
-    const handleRunClick = () => {
+    const {close, error, getPersistedJobId, persistJobId, setStreamRequest} = useWorkflowTestStream(
+        workflow.id!,
+        () => {
+            if (bottomResizablePanelRef.current && bottomResizablePanelRef.current.getSize() === 0) {
+                bottomResizablePanelRef.current.resize(35);
+            }
+
+            setJobId(null);
+        },
+        () => setJobId(null),
+        (jobId) => setJobId(jobId)
+    );
+
+    const handleRunClick = useCallback(() => {
         setShowBottomPanelOpen(true);
         setWorkflowTestExecution(undefined);
 
@@ -116,27 +132,33 @@ export const useWorkflowBuilderHeader = ({bottomResizablePanelRef, chatTrigger,
                 setWorkflowTestChatPanelOpen(true);
             } else {
                 setWorkflowIsRunning(true);
+                setJobId(null);
+                persistJobId(null);
 
-                workflowTestApi
-                    .testWorkflow({
-                        environmentId: currentEnvironmentId,
-                        id: workflow.id,
-                    })
-                    .then((workflowTestExecution) => {
-                        setWorkflowTestExecution(workflowTestExecution);
-                        setWorkflowIsRunning(false);
-
-                        if (bottomResizablePanelRef.current && bottomResizablePanelRef.current.getSize() === 0) {
-                            bottomResizablePanelRef.current.resize(35);
-                        }
-                    })
-                    .catch(() => {
-                        setWorkflowIsRunning(false);
-                        setWorkflowTestExecution(undefined);
-                    });
+                const request = getTestWorkflowStreamPostRequest({
+                    environmentId: currentEnvironmentId,
+                    id: workflow.id,
+                });
+
+                setStreamRequest(request);
             }
         }
-    };
+    }, [
+        captureProjectWorkflowTested,
+        currentEnvironmentId,
+        bottomResizablePanelRef,
+        chatTrigger,
+        persistJobId,
+        resetMessages,
+        setDataPillPanelOpen,
+        setShowBottomPanelOpen,
+        setStreamRequest,
+        setWorkflowIsRunning,
+        setWorkflowNodeDetailsPanelOpen,
+        setWorkflowTestExecution,
+        setWorkflowTestChatPanelOpen,
+        workflow.id,
+    ]);
 
     const handleShowOutputClick = () => {
         setShowBottomPanelOpen(!showBottomPanel);
@@ -148,6 +170,15 @@ export const useWorkflowBuilderHeader = ({bottomResizablePanelRef, chatTrigger,
 
     const handleStopClick = useCallback(() => {
         setWorkflowIsRunning(false);
+        close();
+        setStreamRequest(null);
+
+        if (jobId) {
+            workflowTestApi.stopWorkflowTest({jobId}).finally(() => {
+                persistJobId(null);
+                setJobId(null);
+            });
+        }
 
         if (chatTrigger) {
             setWorkflowTestChatPanelOpen(false);
@@ -156,13 +187,52 @@ export const useWorkflowBuilderHeader = ({bottomResizablePanelRef, chatTrigger,
                 bottomResizablePanelRef.current.resize(0);
             }
         }
-    }, [bottomResizablePanelRef, chatTrigger, setWorkflowIsRunning, setWorkflowTestChatPanelOpen]);
+    }, [
+        bottomResizablePanelRef,
+        chatTrigger,
+        close,
+        jobId,
+        persistJobId,
+        setStreamRequest,
+        setWorkflowIsRunning,
+        setWorkflowTestChatPanelOpen,
+    ]);
+
+    // On mount: try to restore an ongoing run using jobId persisted in localStorage.
+    // Attach-first approach: immediately call attach with the exact jobId string.
+    useEffect(() => {
+        if (!workflow.id || currentEnvironmentId === undefined) return;
+
+        const jobId = getPersistedJobId();
+
+        if (!jobId) {
+            return;
+        }
+
+        setWorkflowIsRunning(true);
+        setJobId(jobId);
+
+        setStreamRequest(getTestWorkflowAttachRequest({jobId}));
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [workflow.id, currentEnvironmentId, getPersistedJobId]);
 
     useEffect(() => {
-        if (workflowNodeDetailsPanelOpen || !workflowTestChatPanelOpen) {
+        // Stop only when:
+        // - Node details panel is opened (always cancels runs), or
+        // - We are in chat mode and the chat panel is not open
+        if (workflowNodeDetailsPanelOpen || (chatTrigger && !workflowTestChatPanelOpen)) {
             handleStopClick();
         }
-    }, [handleStopClick, workflowNodeDetailsPanelOpen, workflowTestChatPanelOpen]);
+    }, [chatTrigger, handleStopClick, workflowNodeDetailsPanelOpen, workflowTestChatPanelOpen]);
+
+    useEffect(() => {
+        if (error) {
+            setWorkflowIsRunning(false);
+            setStreamRequest(null);
+            persistJobId(null);
+            setJobId(null);
+        }
+    }, [error, persistJobId, setWorkflowIsRunning, setStreamRequest]);
 
     return {
         handleProjectWorkflowValueChange,

client/src/pages/automation/project/Project.tsx

@@ -8,31 +8,31 @@ import WorkflowExecutionsTestOutput from '@/pages/platform/workflow-editor/compo
 import {useRun} from '@/pages/platform/workflow-editor/hooks/useRun';
 import {WorkflowEditorProvider} from '@/pages/platform/workflow-editor/providers/workflowEditorProvider';
 import useWorkflowDataStore from '@/pages/platform/workflow-editor/stores/useWorkflowDataStore';
-import useWorkflowEditorStore from '@/pages/platform/workflow-editor/stores/useWorkflowEditorStore';
+import WorkflowTestRunLeaveDialog from '@/shared/components/WorkflowTestRunLeaveDialog';
+import {useWorkflowTestRunGuard} from '@/shared/hooks/useWorkflowTestRunGuard';
 import {WebhookTriggerTestApi} from '@/shared/middleware/automation/configuration';
 import {useCreateConnectionMutation} from '@/shared/mutations/automation/connections.mutations';
 import {useGetComponentDefinitionsQuery} from '@/shared/queries/automation/componentDefinitions.queries';
 import {ConnectionKeys, useGetConnectionTagsQuery} from '@/shared/queries/automation/connections.queries';
+import {useEnvironmentStore} from '@/shared/stores/useEnvironmentStore';
 import {useShallow} from 'zustand/react/shallow';
 
 const Project = () => {
+    const currentEnvironmentId = useEnvironmentStore((state) => state.currentEnvironmentId);
     const {projectLeftSidebarOpen} = useProjectsLeftSidebarStore(
         useShallow((state) => ({
             projectLeftSidebarOpen: state.projectLeftSidebarOpen,
         }))
     );
-    const {workflowIsRunning, workflowTestExecution} = useWorkflowEditorStore(
-        useShallow((state) => ({
-            workflowIsRunning: state.workflowIsRunning,
-            workflowTestExecution: state.workflowTestExecution,
-        }))
-    );
     const {workflow} = useWorkflowDataStore(
         useShallow((state) => ({
             workflow: state.workflow,
         }))
     );
 
+    const {cancelLeave, confirmLeave, showLeaveDialog, workflowIsRunning, workflowTestExecution} =
+        useWorkflowTestRunGuard(workflow.id, currentEnvironmentId);
+
     const {
         bottomResizablePanelRef,
         deleteClusterElementParameterMutation,
@@ -49,11 +49,12 @@ const Project = () => {
         updateWorkflowNodeParameterMutation,
         useGetConnectionsQuery,
     } = useProject();
-
     const {runDisabled} = useRun();
 
     return (
         <div className="flex w-full">
+            <WorkflowTestRunLeaveDialog onCancel={cancelLeave} onConfirm={confirmLeave} open={showLeaveDialog} />
+
             {projectLeftSidebarOpen && projects && (
                 <ProjectsLeftSidebar
                     bottomResizablePanelRef={bottomResizablePanelRef}