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

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,51 @@ 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}));
+    }, [workflow.id, currentEnvironmentId, getPersistedJobId, setWorkflowIsRunning, setJobId, setStreamRequest]);
 
     useEffect(() => {
-        if (workflowNodeDetailsPanelOpen || !workflowTestChatPanelOpen) {
+        // Stop only when:
+        // - The 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/ee/pages/embedded/integration/Integration.tsx

@@ -16,27 +16,33 @@ import WorkflowEditorLayout from '@/pages/platform/workflow-editor/WorkflowEdito
 import WorkflowExecutionsTestOutput from '@/pages/platform/workflow-editor/components/WorkflowExecutionsTestOutput';
 import {useRun} from '@/pages/platform/workflow-editor/hooks/useRun';
 import {WorkflowEditorProvider} from '@/pages/platform/workflow-editor/providers/workflowEditorProvider';
-import useWorkflowEditorStore from '@/pages/platform/workflow-editor/stores/useWorkflowEditorStore';
+import useWorkflowDataStore from '@/pages/platform/workflow-editor/stores/useWorkflowDataStore';
+import WorkflowTestRunLeaveDialog from '@/shared/components/WorkflowTestRunLeaveDialog';
+import {useWorkflowTestRunGuard} from '@/shared/hooks/useWorkflowTestRunGuard';
 import Header from '@/shared/layout/Header';
 import LayoutContainer from '@/shared/layout/LayoutContainer';
 import {WebhookTriggerTestApi} from '@/shared/middleware/automation/configuration';
+import {useEnvironmentStore} from '@/shared/stores/useEnvironmentStore';
 import {useQueryClient} from '@tanstack/react-query';
 import {useParams} from 'react-router-dom';
 import {useShallow} from 'zustand/react/shallow';
 
 const Integration = () => {
+    const currentEnvironmentId = useEnvironmentStore((state) => state.currentEnvironmentId);
     const {leftSidebarOpen} = useIntegrationsLeftSidebarStore(
         useShallow((state) => ({
             leftSidebarOpen: state.leftSidebarOpen,
         }))
     );
-    const {workflowIsRunning, workflowTestExecution} = useWorkflowEditorStore(
+    const {workflow} = useWorkflowDataStore(
         useShallow((state) => ({
-            workflowIsRunning: state.workflowIsRunning,
-            workflowTestExecution: state.workflowTestExecution,
+            workflow: state.workflow,
         }))
     );
 
+    const {cancelLeave, confirmLeave, showLeaveDialog, workflowIsRunning, workflowTestExecution} =
+        useWorkflowTestRunGuard(workflow.id, currentEnvironmentId);
+
     const {integrationId, integrationWorkflowId} = useParams();
 
     const queryClient = useQueryClient();
@@ -59,6 +65,7 @@ const Integration = () => {
 
     return (
         <>
+            <WorkflowTestRunLeaveDialog onCancel={cancelLeave} onConfirm={confirmLeave} open={showLeaveDialog} />
             <LayoutContainer
                 className="bg-muted/50"
                 leftSidebarBody={<IntegrationsSidebar integrationId={+integrationId!} />}

client/src/ee/pages/embedded/integration/components/integration-header/IntegrationHeader.tsx

@@ -29,12 +29,14 @@ import useWorkflowNodeDetailsPanelStore from '@/pages/platform/workflow-editor/s
 import DeleteWorkflowAlertDialog from '@/shared/components/DeleteWorkflowAlertDialog';
 import WorkflowDialog from '@/shared/components/workflow/WorkflowDialog';
 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 {UpdateWorkflowMutationType} from '@/shared/types';
+import {getTestWorkflowAttachRequest, getTestWorkflowStreamPostRequest} from '@/shared/util/testWorkflow-utils';
 import {useQueryClient} from '@tanstack/react-query';
 import {PlusIcon} from 'lucide-react';
-import {RefObject, useState} from 'react';
+import {RefObject, useCallback, useEffect, useState} from 'react';
 import {ImperativePanelHandle} from 'react-resizable-panels';
 import {useLoaderData, useNavigate, useSearchParams} from 'react-router-dom';
 import {useShallow} from 'zustand/react/shallow';
@@ -54,6 +56,7 @@ const IntegrationHeader = ({
     runDisabled: boolean;
     updateWorkflowMutation: UpdateWorkflowMutationType;
 }) => {
+    const [jobId, setJobId] = useState<string | null>(null);
     const [showDeleteIntegrationAlertDialog, setShowDeleteIntegrationAlertDialog] = useState(false);
     const [showDeleteWorkflowAlertDialog, setShowDeleteWorkflowAlertDialog] = useState(false);
     const [showEditIntegrationDialog, setShowEditIntegrationDialog] = useState(false);
@@ -82,6 +85,8 @@ const IntegrationHeader = ({
 
     const {captureIntegrationWorkflowCreated, captureIntegrationWorkflowTested} = useAnalytics();
 
+    const queryClient = useQueryClient();
+
     const navigate = useNavigate();
 
     const [searchParams] = useSearchParams();
@@ -92,7 +97,18 @@ const IntegrationHeader = ({
         !showDeleteIntegrationAlertDialog
     );
 
-    const queryClient = useQueryClient();
+    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 createIntegrationWorkflowMutation = useCreateIntegrationWorkflowMutation({
         onSuccess: (integrationWorkflowId) => {
@@ -173,13 +189,12 @@ const IntegrationHeader = ({
         setCurrentNode(undefined);
 
         navigate(
-            `/embedded/integrations/${integrationId}/integration-workflows/${integrationWorkflowId}?${searchParams}`
+            `/embedded/integrations/${integrationId}/integration-workflows/${integrationWorkflowId}?${searchParams.toString()}`
         );
     };
 
-    const handleRunClick = () => {
+    const handleRunClick = useCallback(() => {
         setShowBottomPanelOpen(true);
-        setWorkflowIsRunning(true);
         setWorkflowTestExecution(undefined);
 
         if (bottomResizablePanelRef.current) {
@@ -189,25 +204,65 @@ const IntegrationHeader = ({
         if (workflow.id) {
             captureIntegrationWorkflowTested();
 
-            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);
-                });
+            setWorkflowIsRunning(true);
+            setJobId(null);
+            persistJobId(null);
+
+            const req = getTestWorkflowStreamPostRequest({
+                environmentId: currentEnvironmentId,
+                id: workflow.id,
+            });
+            setStreamRequest(req);
         }
-    };
+    }, [
+        captureIntegrationWorkflowTested,
+        currentEnvironmentId,
+        bottomResizablePanelRef,
+        persistJobId,
+        setShowBottomPanelOpen,
+        setStreamRequest,
+        setWorkflowIsRunning,
+        setWorkflowTestExecution,
+        workflow.id,
+    ]);
+
+    const handleStopClick = useCallback(() => {
+        setWorkflowIsRunning(false);
+        close();
+        setStreamRequest(null);
+
+        if (jobId) {
+            workflowTestApi.stopWorkflowTest({jobId}).finally(() => {
+                persistJobId(null);
+                setJobId(null);
+            });
+        }
+    }, [close, jobId, persistJobId, setStreamRequest, setWorkflowIsRunning]);
+
+    useEffect(() => {
+        if (!workflow.id || currentEnvironmentId === undefined) return;
+
+        const jobId = getPersistedJobId();
+
+        if (!jobId) {
+            return;
+        }
+
+        setWorkflowIsRunning(true);
+        setJobId(jobId);
+
+        setStreamRequest(getTestWorkflowAttachRequest({jobId}));
+    }, [workflow.id, currentEnvironmentId, getPersistedJobId, setWorkflowIsRunning, setJobId, setStreamRequest]);
+
+    // On transport error (e.g., 4xx/5xx), make sure to reset running state and clear the request to prevent retries
+    useEffect(() => {
+        if (error) {
+            setWorkflowIsRunning(false);
+            setStreamRequest(null);
+            persistJobId(null);
+            setJobId(null);
+        }
+    }, [error, persistJobId, setWorkflowIsRunning, setStreamRequest]);
 
     return (
         <header className="flex items-center px-3 py-2.5">
@@ -246,7 +301,7 @@ const IntegrationHeader = ({
                     )}
 
                     {workflowIsRunning ? (
-                        <IntegrationHeaderStopButton />
+                        <IntegrationHeaderStopButton onClick={handleStopClick} />
                     ) : (
                         <IntegrationHeaderRunButton onRunClick={handleRunClick} runDisabled={runDisabled} />
                     )}

client/src/ee/pages/embedded/integration/components/integration-header/IntegrationHeaderStopButton.tsx

@@ -1,15 +1,8 @@
 import {Button} from '@/components/ui/button';
 import {SquareIcon} from 'lucide-react';
 
-const IntegrationHeaderStopButton = () => (
-    <Button
-        className="hover:bg-background/70 [&_svg]:size-5"
-        onClick={() => {
-            // TODO
-        }}
-        size="icon"
-        variant="ghost"
-    >
+const IntegrationHeaderStopButton = ({onClick}: {onClick?: () => void}) => (
+    <Button className="hover:bg-background/70 [&_svg]:size-5" onClick={onClick} size="icon" variant="ghost">
         <SquareIcon className="text-destructive" />
     </Button>
 );