fix: wf ai assistant more accurate steps

keep-ui/entities/workflows/model/__tests__/workflow-store.test.tsx

@@ -66,7 +66,7 @@ describe("useWorkflowStore", () => {
 
       // Add a trigger node
       act(() => {
-        result.current.addNodeBetween(
+        result.current.addNodeBetweenSafe(
           "edge-1",
           {
             id: "interval",
@@ -109,7 +109,7 @@ describe("useWorkflowStore", () => {
       // Try to add another trigger
       act(() => {
         const edges = result.current.edges;
-        result.current.addNodeBetween(
+        result.current.addNodeBetweenSafe(
           edges[1].id,
           {
             id: "interval",

keep-ui/entities/workflows/model/types.ts

@@ -170,6 +170,13 @@ export type V2StepConditionThreshold = z.infer<
   typeof V2StepConditionThresholdSchema
 >;
 
+export const V2StepConditionSchema = z.union([
+  V2StepConditionAssertSchema,
+  V2StepConditionThresholdSchema,
+]);
+
+export type V2StepCondition = z.infer<typeof V2StepConditionSchema>;
+
 export const V2StepForeachSchema = z.object({
   id: z.string(),
   name: z.string(),
@@ -347,6 +354,11 @@ export interface FlowState extends FlowStateValues {
     step: V2StepTemplate | V2StepTrigger,
     type: "node" | "edge"
   ) => string | null;
+  addNodeBetweenSafe: (
+    nodeOrEdgeId: string,
+    step: V2StepTemplate | V2StepTrigger,
+    type: "node" | "edge"
+  ) => string | null;
   setToolBoxConfig: (config: ToolboxConfiguration) => void;
   setEditorOpen: (open: boolean) => void;
   updateSelectedNodeData: (key: string, value: any) => void;
@@ -361,7 +373,7 @@ export interface FlowState extends FlowStateValues {
   setEdges: (edges: Edge[]) => void;
   getNodeById: (id: string) => FlowNode | undefined;
   getEdgeById: (id: string) => Edge | undefined;
-  deleteNodes: (ids: string | string[]) => void;
+  deleteNodes: (ids: string | string[]) => string[];
   getNextEdge: (nodeId: string) => Edge | undefined;
   reset: () => void;
   setDefinition: (def: DefinitionV2) => void;

keep-ui/entities/workflows/model/workflow-store.ts

@@ -260,6 +260,15 @@ export const useWorkflowStore = create<FlowState>()(
       nodeOrEdgeId: string,
       step: V2StepTemplate | V2StepTrigger,
       type: "node" | "edge"
+    ) => {
+      const newNodeId = addNodeBetween(nodeOrEdgeId, step, type, set, get);
+      set({ selectedNode: newNodeId, selectedEdge: null });
+      return newNodeId ?? null;
+    },
+    addNodeBetweenSafe: (
+      nodeOrEdgeId: string,
+      step: V2StepTemplate | V2StepTrigger,
+      type: "node" | "edge"
     ) => {
       try {
         const newNodeId = addNodeBetween(nodeOrEdgeId, step, type, set, get);
@@ -474,12 +483,12 @@ export const useWorkflowStore = create<FlowState>()(
     deleteNodes: (ids) => {
       //for now handling only single node deletion. can later enhance to multiple deletions
       if (typeof ids !== "string") {
-        return;
+        return [];
       }
       const nodes = get().nodes;
       const nodeStartIndex = nodes.findIndex((node) => ids == node.id);
       if (nodeStartIndex === -1) {
-        return;
+        return [];
       }
       let idArray = Array.isArray(ids) ? ids : [ids];
 
@@ -561,6 +570,8 @@ export const useWorkflowStore = create<FlowState>()(
       });
       get().onLayout({ direction: "DOWN" });
       get().updateDefinition();
+
+      return [ids];
     },
     getNextEdge: (nodeId: string) => {
       const node = get().getNodeById(nodeId);

keep-ui/features/workflows/builder/ui/BuilderChat/AddStepUI.tsx

@@ -1,157 +1,64 @@
 import { Button } from "@/components/ui";
-import { useWorkflowStore, V2Step } from "@/entities/workflows";
-import { WF_DEBUG_INFO } from "../debug-settings";
-import { DebugArgs } from "./debug-args";
 import { StepPreview } from "./StepPreview";
 import { SuggestionResult, SuggestionStatus } from "./SuggestionStatus";
 import clsx from "clsx";
-import { DebugJSON } from "@/shared/ui";
-import { useCallback } from "react";
-import { triggerTypes } from "../../lib/utils";
-import { Edge } from "@xyflow/react";
+import { V2Step } from "@/entities/workflows/model/types";
+import { useWorkflowStore } from "@/entities/workflows";
 
-type AddStepUIProps =
-  | {
-      status: "complete";
-      args: {
-        stepDefinitionJSON?: string;
-        addAfterNodeName?: string;
-        addAfterEdgeId?: string;
-        isStart?: boolean;
-      };
-      respond: undefined;
-      result: SuggestionResult;
-    }
-  | {
-      status: "executing";
-      args: {
-        stepDefinitionJSON?: string;
-        addAfterNodeName?: string;
-        addAfterEdgeId?: string;
-        isStart?: boolean;
-      };
-      respond: (response: SuggestionResult) => void;
-      result: undefined;
-    };
+type AddStepUIPropsCommon = {
+  step: V2Step;
+  addAfterEdgeId: string;
+};
 
-function getComponentType(stepType: string) {
-  if (stepType.startsWith("step-")) return "task";
-  if (stepType.startsWith("action-")) return "task";
-  if (stepType.startsWith("condition-")) return "switch";
-  if (stepType === "foreach") return "container";
-  if (triggerTypes.includes(stepType)) return "trigger";
-  return "task";
-}
+type AddStepUIPropsComplete = AddStepUIPropsCommon & {
+  status: "complete";
+  result: SuggestionResult;
+  respond: undefined;
+};
 
-function parseAndAutoCorrectStepDefinition(stepDefinitionJSON: string) {
-  const step = JSON.parse(stepDefinitionJSON);
-  return {
-    ...step,
-    componentType: getComponentType(step.type),
-  };
-}
+type AddStepUIPropsExecuting = AddStepUIPropsCommon & {
+  status: "executing";
+  result: undefined;
+  respond: (response: SuggestionResult) => void;
+};
+
+type AddStepUIProps = AddStepUIPropsComplete | AddStepUIPropsExecuting;
 
 export const AddStepUI = ({
   status,
-  args,
-  respond,
+  step,
+  addAfterEdgeId,
   result,
+  respond,
 }: AddStepUIProps) => {
-  const {
-    definition,
-    nodes,
-    getNodeById,
-    getNextEdge,
-    addNodeBetween,
-    getEdgeById,
-  } = useWorkflowStore();
-  let {
-    stepDefinitionJSON,
-    addAfterNodeName: addAfterNodeIdOrName,
-    addAfterEdgeId,
-    isStart,
-  } = args;
+  const { addNodeBetween } = useWorkflowStore();
 
-  const addNodeAfterNode = useCallback(
-    (
-      nodeToAddAfterId: string,
-      step: V2Step,
-      isStart: boolean,
-      respond: (response: any) => void,
-      addAfterEdgeId?: string
-    ) => {
-      if (
-        nodeToAddAfterId === "alert" ||
-        nodeToAddAfterId === "incident" ||
-        nodeToAddAfterId === "interval" ||
-        nodeToAddAfterId === "manual"
-      ) {
-        nodeToAddAfterId = "trigger_end";
-      }
-      let node = getNodeById(isStart ? "trigger_end" : nodeToAddAfterId);
-      if (!node) {
-        const nodeByName = definition?.value.sequence.find(
-          (s) => s.name === nodeToAddAfterId
-        );
-        if (nodeByName) {
-          node = getNodeById(nodeByName.id);
-        }
-        if (!node) {
-          respond?.({
-            status: "error",
-            error: new Error("Can't find the node to add the step after"),
-            message: "Step not added due to error",
-          });
-          return;
-        }
-      }
-      let nextEdge: Edge | null = null;
-      if (addAfterEdgeId) {
-        nextEdge = getEdgeById(addAfterEdgeId) ?? null;
-      }
-      if (!nextEdge) {
-        nextEdge = getNextEdge(node.id) ?? null;
-      }
-      if (!nextEdge) {
-        respond?.({
-          status: "error",
-          error: new Error("Can't find the edge to add the step after"),
-          message: "Step not added due to error",
-        });
-        return;
-      }
-      try {
-        addNodeBetween(nextEdge.id, step, "edge");
-        respond?.({
-          status: "complete",
-          stepId: step.id,
-          message: "Step added",
-        });
-      } catch (e) {
-        respond?.({
-          status: "error",
-          error: e,
-          message: "Step not added due to error",
-        });
-      }
-    },
-    [addNodeBetween, definition?.value.sequence, getNextEdge, getNodeById]
-  );
+  const onAdd = () => {
+    try {
+      addNodeBetween(addAfterEdgeId, step, "edge");
+      respond?.({
+        status: "complete",
+        message: "Step added",
+      });
+    } catch (e) {
+      respond?.({
+        status: "error",
+        error: e,
+        message: "Step not added",
+      });
+    }
+  };
 
-  if (!stepDefinitionJSON) {
-    return <div>Step definition not found</div>;
-  }
-  if (definition?.value.sequence.length === 0) {
-    isStart = true;
-  }
-  let step = parseAndAutoCorrectStepDefinition(stepDefinitionJSON);
+  const onCancel = () => {
+    respond?.({
+      status: "complete",
+      message: "User cancelled adding step",
+    });
+  };
 
   if (status === "complete") {
     return (
       <div className="flex flex-col gap-1 my-2">
-        {WF_DEBUG_INFO && (
-          <DebugArgs args={{ isStart, addAfterNodeIdOrName }} nodes={nodes} />
-        )}
         <StepPreview
           step={step}
           className={clsx(
@@ -167,58 +74,17 @@ export const AddStepUI = ({
   return (
     <div className="flex flex-col gap-2">
       <div>
-        <div>
-          Do you want to add this step after <b>{addAfterNodeIdOrName}</b>
-          <pre>{step.name}</pre>
-          {WF_DEBUG_INFO && (
-            <DebugArgs args={{ isStart, addAfterNodeIdOrName }} nodes={nodes} />
-          )}
-          {WF_DEBUG_INFO && (
-            <DebugJSON
-              name="stepDefinitionJSON"
-              json={JSON.parse(stepDefinitionJSON ?? "")}
-            />
-          )}
-        </div>
+        {/* TODO: add the place where the action will be added in text */}
+        <div>Do you want to add this action?</div>
         <div className="my-2">
           <StepPreview step={step} />
         </div>
       </div>
       <div className="flex gap-2">
-        <Button
-          color="orange"
-          variant="primary"
-          onClick={async () => {
-            try {
-              addNodeAfterNode(
-                addAfterNodeIdOrName ?? "",
-                step,
-                !!isStart,
-                respond,
-                addAfterEdgeId
-              );
-            } catch (e) {
-              console.error(e);
-              respond?.({
-                status: "error",
-                error: e,
-                message: `Error adding step: ${e}`,
-              });
-            }
-          }}
-        >
+        <Button color="orange" variant="primary" onClick={onAdd}>
           Add (⌘+Enter)
         </Button>
-        <Button
-          color="orange"
-          variant="secondary"
-          onClick={() =>
-            respond?.({
-              status: "declined",
-              message: "Step suggestion declined",
-            })
-          }
-        >
+        <Button color="orange" variant="secondary" onClick={onCancel}>
           No
         </Button>
       </div>

keep-ui/features/workflows/builder/ui/BuilderChat/AddTriggerUI.tsx

@@ -56,7 +56,11 @@ export const AddTriggerUI = ({
   result,
 }: AddTriggerUIProps) => {
   const [isAddingTrigger, setIsAddingTrigger] = useState(false);
-  const { nodes, addNodeBetween, getNextEdge } = useWorkflowStore();
+  const {
+    nodes,
+    addNodeBetweenSafe: addNodeBetween,
+    getNextEdge,
+  } = useWorkflowStore();
   const { triggerType, triggerProperties } = args;
 
   const triggerDefinition = useMemo(() => {