From b62ca684e0396873ae443453c218816255f80614 Mon Sep 17 00:00:00 2001
From: Lazaros Toumanidis <laztoum@protonmail.com>
Date: Thu, 16 Jan 2025 21:39:19 +0200
Subject: [PATCH] update swarm edge styles, connect user to swarm, initial
 agent

---
 src/waldiez/containers/edges/main.tsx         |   2 +-
 src/waldiez/containers/edges/swarm.tsx        |  33 ++--
 src/waldiez/containers/edges/utils.ts         |  67 +++++++
 src/waldiez/containers/nodes/agent/main.tsx   | 178 ++++++++----------
 .../nodes/agent/swarmContainer/modal/main.tsx |   6 +-
 .../editFlowModal/hooks/useEditFlowModal.ts   |   4 +-
 src/waldiez/models/Stores/IAgentStore.ts      |   2 +
 src/waldiez/models/Stores/IFlowStore.ts       |   2 +-
 src/waldiez/store/agent.ts                    |  42 ++++-
 src/waldiez/store/edge.ts                     |  21 ++-
 src/waldiez/store/flow.ts                     |  12 +-
 src/waldiez/store/utils/agent.ts              |  36 ++++
 src/waldiez/store/utils/edge.ts               |  64 +++----
 src/waldiez/styles/overrides.css              |  80 +-------
 src/waldiez/styles/swarm.css                  |  23 +++
 15 files changed, 318 insertions(+), 254 deletions(-)
 create mode 100644 src/waldiez/containers/edges/utils.ts

diff --git a/src/waldiez/containers/edges/main.tsx b/src/waldiez/containers/edges/main.tsx
index 9fd2550..540584a 100644
--- a/src/waldiez/containers/edges/main.tsx
+++ b/src/waldiez/containers/edges/main.tsx
@@ -75,7 +75,7 @@ const WaldiezEdgeCommon = (props: WaldiezEdgeProps) => {
             return <WaldiezEdgeSwarmView {...props} swarmType="nested" />;
         }
         const isSwarmSource = () =>
-            sourceAgent.data.agentType !== "swarm" && targetAgent?.data.agentType === "swarm_container";
+            sourceAgent.data.agentType !== "swarm" && targetAgent?.data.agentType === "swarm";
         if (isSwarmSource()) {
             return <WaldiezEdgeSwarmView {...props} swarmType="source" />;
         }
diff --git a/src/waldiez/containers/edges/swarm.tsx b/src/waldiez/containers/edges/swarm.tsx
index 184af59..3f1b606 100644
--- a/src/waldiez/containers/edges/swarm.tsx
+++ b/src/waldiez/containers/edges/swarm.tsx
@@ -1,4 +1,4 @@
-import { BaseEdge, EdgeLabelRenderer, Position, getSmoothStepPath } from "@xyflow/react";
+import { BaseEdge, EdgeLabelRenderer, getSmoothStepPath } from "@xyflow/react";
 
 import { FaTrashAlt } from "react-icons/fa";
 import { FaGear } from "react-icons/fa6";
@@ -7,6 +7,7 @@ import { GiShakingHands } from "react-icons/gi";
 import { MdMessage } from "react-icons/md";
 
 import { WaldiezEdgeProps } from "@waldiez/containers/edges/types";
+import { getEdgeTranslations } from "@waldiez/containers/edges/utils";
 import { useWaldiez } from "@waldiez/store";
 import { AGENT_COLORS } from "@waldiez/theme";
 
@@ -118,7 +119,7 @@ export const WaldiezEdgeSwarmView = (
         if (label === "") {
             return null;
         }
-        const trimmedTo20 = label.length > 15 ? `${label.slice(0, 15)}...` : label;
+        const trimmedTo20 = label.length > 20 ? `${label.slice(0, 20)}...` : label;
         return (
             <div
                 style={{
@@ -138,26 +139,14 @@ export const WaldiezEdgeSwarmView = (
         );
     };
     const className = swarmType === "source" ? "agent-edge-box" : "clickable agent-edge-swarm-box";
-    const translations = {
-        edgeStart: `translate(-50%, 0%) translate(${sourceX}px,${sourceY}px)`,
-        edgeEnd: `translate(-50%, 0%) translate(${targetX}px,${targetY}px)`,
-    };
-    if (sourcePosition === Position.Right && targetPosition === Position.Left) {
-        translations.edgeStart = `translate(0%, 0%) translate(${sourceX}px,${sourceY - 35}px)`;
-        translations.edgeEnd = `translate(-100%, -100%) translate(${targetX}px,${targetY}px)`;
-    }
-    if (sourcePosition === Position.Left && targetPosition === Position.Right) {
-        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX}px,${sourceY}px)`;
-        translations.edgeEnd = `translate(0, 0) translate(${targetX}px,${targetY}px)`;
-    }
-    if (sourcePosition === Position.Top && targetPosition === Position.Bottom) {
-        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX}px,${sourceY - 30}px)`;
-        translations.edgeEnd = `translate(-100%, 0%) translate(${targetX}px,${targetY}px)`;
-    }
-    if (sourcePosition === Position.Bottom && targetPosition === Position.Top) {
-        translations.edgeStart = `translate(0%, 0%) translate(${sourceX}px,${sourceY}px)`;
-        translations.edgeEnd = `translate(0%, 0%) translate(${targetX}px,${targetY - 30}px)`;
-    }
+    const translations = getEdgeTranslations(
+        sourceX,
+        sourceY,
+        targetX,
+        targetY,
+        sourcePosition,
+        targetPosition,
+    );
     return (
         <>
             <BaseEdge path={edgePath} markerEnd={markerEnd} style={style} />
diff --git a/src/waldiez/containers/edges/utils.ts b/src/waldiez/containers/edges/utils.ts
new file mode 100644
index 0000000..7122860
--- /dev/null
+++ b/src/waldiez/containers/edges/utils.ts
@@ -0,0 +1,67 @@
+import { Position } from "@xyflow/react";
+
+// eslint-disable-next-line complexity
+export const getEdgeTranslations = (
+    sourceX: number,
+    sourceY: number,
+    targetX: number,
+    targetY: number,
+    sourcePosition: Position,
+    targetPosition: Position,
+) => {
+    const translations = {
+        edgeStart: `translate(-50%, 0%) translate(${sourceX}px,${sourceY}px)`,
+        edgeEnd: `translate(-50%, 0%) translate(${targetX}px,${targetY}px)`,
+    };
+    if (sourcePosition === Position.Right && targetPosition === Position.Left) {
+        translations.edgeStart = `translate(0%, 0%) translate(${sourceX - 10}px,${sourceY - 35}px)`;
+        translations.edgeEnd = `translate(-100%, -100%) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Right && targetPosition === Position.Top) {
+        translations.edgeStart = `translate(0%, 0%) translate(${sourceX - 10}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(-100%, 0%) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Right && targetPosition === Position.Right) {
+        translations.edgeStart = `translate(0%, 0%) translate(${sourceX - 10}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(0, 0) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Left && targetPosition === Position.Right) {
+        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX + 10}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(0, 0) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Left && targetPosition === Position.Bottom) {
+        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX + 10}px,${sourceY - 30}px)`;
+        translations.edgeEnd = `translate(0, 0) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Left && targetPosition === Position.Left) {
+        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX + 10}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(-100%, 0%) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Top && targetPosition === Position.Bottom) {
+        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX}px,${sourceY - 30}px)`;
+        translations.edgeEnd = `translate(-100%, 0%) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Top && targetPosition === Position.Right) {
+        translations.edgeStart = `translate(-100%, 0%) translate(${sourceX}px,${sourceY - 30}px)`;
+        translations.edgeEnd = `translate(0, 0) translate(${targetX}px,${targetY}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Bottom && targetPosition === Position.Left) {
+        translations.edgeStart = `translate(0%, 0%) translate(${sourceX}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(-100%, 0%) translate(${targetX}px,${targetY - 30}px)`;
+        return translations;
+    }
+    if (sourcePosition === Position.Bottom && targetPosition === Position.Top) {
+        translations.edgeStart = `translate(0%, 0%) translate(${sourceX}px,${sourceY}px)`;
+        translations.edgeEnd = `translate(0%, 0%) translate(${targetX}px,${targetY - 30}px)`;
+        return translations;
+    }
+    return translations;
+};
diff --git a/src/waldiez/containers/nodes/agent/main.tsx b/src/waldiez/containers/nodes/agent/main.tsx
index 56efddd..4874ae6 100644
--- a/src/waldiez/containers/nodes/agent/main.tsx
+++ b/src/waldiez/containers/nodes/agent/main.tsx
@@ -39,6 +39,7 @@ export const WaldiezNodeAgentView = (props: WaldiezNodeAgentProps) => {
             />
         );
     }
+    const handleClassNameBase = agentType === "swarm" ? "swarm-" : data.parentId ? "hidden" : "";
     return (
         <div className={className} data-testid={`agent-node-${id}-view`}>
             {!data.parentId && (
@@ -57,103 +58,86 @@ export const WaldiezNodeAgentView = (props: WaldiezNodeAgentProps) => {
                     <WaldiezNodeAgentFooter id={id} data={data} isModalOpen={isModalOpen} />
                 </div>
             )}
-            {agentType === "swarm" ? (
-                <>
-                    <Handle
-                        className="swarm-handle top target"
-                        type="target"
-                        isConnectableStart
-                        position={Position.Top}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-top-target`}
-                        id={`agent-handle-${id}-top-target`}
-                    />
-                    <Handle
-                        className="swarm-handle top source"
-                        type="source"
-                        isConnectableStart
-                        position={Position.Top}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-top-source`}
-                        id={`agent-handle-${id}-top-source`}
-                    />
-                    <Handle
-                        className="swarm-handle bottom target"
-                        type="target"
-                        isConnectableStart
-                        position={Position.Bottom}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-bottom-target`}
-                        id={`agent-handle-${id}-bottom-target`}
-                    />
-                    <Handle
-                        className="swarm-handle bottom source"
-                        type="source"
-                        isConnectableStart
-                        position={Position.Bottom}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-bottom-source`}
-                        id={`agent-handle-${id}-bottom-source`}
-                    />
-                    <Handle
-                        className="swarm-handle left target"
-                        type="target"
-                        isConnectableStart
-                        position={Position.Left}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-left-target`}
-                        id={`agent-handle-${id}-left-target`}
-                    />
-                    <Handle
-                        className="swarm-handle left source"
-                        type="source"
-                        isConnectableStart
-                        position={Position.Left}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-left-source`}
-                        id={`agent-handle-${id}-left-target`}
-                    />
-                    <Handle
-                        className="swarm-handle right target"
-                        type="target"
-                        isConnectableStart
-                        position={Position.Right}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-right-target`}
-                        id={`agent-handle-${id}-right-target`}
-                    />
-                    <Handle
-                        className="swarm-handle right source"
-                        type="source"
-                        isConnectableStart
-                        position={Position.Right}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-right-source`}
-                        id={`agent-handle-${id}-right-source`}
-                    />
-                </>
-            ) : (
-                <>
-                    <Handle
-                        className={data.parentId ? "hidden" : ""}
-                        type="target"
-                        isConnectableEnd
-                        position={Position.Left}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-left"}`}
-                        id={`agent-handle-${id}-left`}
-                    />
-                    <Handle
-                        className={data.parentId ? "hidden" : ""}
-                        type="source"
-                        isConnectableStart
-                        position={Position.Right}
-                        onConnect={onEdgeConnection}
-                        data-testid={`agent-handle-${id}-right`}
-                        id={`agent-handle-${id}-source`}
-                    />
-                </>
-            )}
+            <Handle
+                className={`${handleClassNameBase}handle top target`}
+                type="target"
+                isConnectableEnd
+                position={Position.Top}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-top-target-${id}`}
+                id={`agent-handle-top-target-${id}`}
+                style={{ left: "75%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle top source`}
+                type="source"
+                isConnectableStart
+                position={Position.Top}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-top-source-${id}`}
+                id={`agent-handle-top-source-${id}`}
+                style={{ left: "25%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle bottom target`}
+                type="target"
+                isConnectableEnd
+                position={Position.Bottom}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-bottom-target-${id}`}
+                id={`agent-handle-bottom-target-${id}`}
+                style={{ left: "25%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle bottom source`}
+                type="source"
+                isConnectableStart
+                position={Position.Bottom}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-bottom-source${id}`}
+                id={`agent-handle-bottom-source${id}`}
+                style={{ left: "75%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle left target`}
+                type="target"
+                // isConnectableEnd
+                position={Position.Left}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-left-target-${id}`}
+                id={`agent-handle-left-target-${id}`}
+                style={{ top: "25%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle left source`}
+                type="source"
+                isConnectableStart
+                position={Position.Left}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-left-source-${id}`}
+                id={`agent-handle-left-source-${id}`}
+                style={{ top: "75%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle right target`}
+                type="target"
+                isConnectableEnd
+                position={Position.Right}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-right-target-${id}`}
+                id={`agent-handle-right-target-${id}`}
+                style={{ top: "75%" }}
+            />
+            <Handle
+                className={`${handleClassNameBase}handle right source`}
+                type="source"
+                isConnectableStart
+                position={Position.Right}
+                onConnect={onEdgeConnection}
+                data-testid={`agent-handle-right-source-${id}`}
+                id={`agent-handle-right-source-${id}`}
+                style={{ top: "25%" }}
+            />
             <button
                 title="Open Node Modal"
                 type="button"
diff --git a/src/waldiez/containers/nodes/agent/swarmContainer/modal/main.tsx b/src/waldiez/containers/nodes/agent/swarmContainer/modal/main.tsx
index 47fe29a..572eb03 100644
--- a/src/waldiez/containers/nodes/agent/swarmContainer/modal/main.tsx
+++ b/src/waldiez/containers/nodes/agent/swarmContainer/modal/main.tsx
@@ -33,6 +33,7 @@ export const WaldiezNodeSwarmContainerModal = (props: WaldiezSwarmContainerModal
     const getSwarmEdges = useWaldiez(s => s.getSwarmEdges);
     const getNonSwarmAgents = useWaldiez(s => s.getNonSwarmAgents);
     const updateAgentData = useWaldiez(s => s.updateAgentData);
+    const updateSwarmInitialAgent = useWaldiez(s => s.updateSwarmInitialAgent);
     const onFlowChanged = useWaldiez(s => s.onFlowChanged);
     const onSave = () => {
         updateAgentData(id, agentData);
@@ -61,7 +62,10 @@ export const WaldiezNodeSwarmContainerModal = (props: WaldiezSwarmContainerModal
     };
     const onInitialAgentChange = (option: SingleValue<{ label: string; value: any }>) => {
         const newData = { ...agentData, initialAgent: option?.value.id };
-        setAgentData(newData);
+        if (option?.value.id) {
+            updateSwarmInitialAgent(option.value.id);
+        }
+        // setAgentData(newData);
         const dirty = !isEqual(data, newData);
         setIsDirty(dirty);
     };
diff --git a/src/waldiez/containers/sidebar/modals/editFlowModal/hooks/useEditFlowModal.ts b/src/waldiez/containers/sidebar/modals/editFlowModal/hooks/useEditFlowModal.ts
index 873a915..632b4be 100644
--- a/src/waldiez/containers/sidebar/modals/editFlowModal/hooks/useEditFlowModal.ts
+++ b/src/waldiez/containers/sidebar/modals/editFlowModal/hooks/useEditFlowModal.ts
@@ -25,7 +25,7 @@ export const useEditFlowModal = (props: EditFlowModalProps) => {
     });
     const [selectedNewEdge, setSelectedNewEdge] = useState<WaldiezEdge | null>(null);
     const getFlowEdges = useWaldiez(s => s.getFlowEdges);
-    const { used: sortedEdges, remaining: remainingEdges } = getFlowEdges();
+    const { used: sortedEdges, remaining: remainingEdges } = getFlowEdges(true);
     const onFlowChanged = useWaldiez(s => s.onFlowChanged);
     // tmp state (to save onSubmit, discard onCancel)
     const [sortedEdgesState, setSortedEdgesState] = useState<WaldiezEdge[]>(sortedEdges);
@@ -56,7 +56,7 @@ export const useEditFlowModal = (props: EditFlowModalProps) => {
     const reset = () => {
         const { name, description, requirements, tags, isAsync } = getFlowInfo();
         setFlowData({ name, description, requirements, tags, isAsync });
-        const { used, remaining } = getFlowEdges();
+        const { used, remaining } = getFlowEdges(true);
         setSortedEdgesState(used);
         setRemainingEdgeState(remaining);
         setIsDirty(false);
diff --git a/src/waldiez/models/Stores/IAgentStore.ts b/src/waldiez/models/Stores/IAgentStore.ts
index 2fdaf3a..2ccc7ad 100644
--- a/src/waldiez/models/Stores/IAgentStore.ts
+++ b/src/waldiez/models/Stores/IAgentStore.ts
@@ -29,6 +29,8 @@ export interface IWaldiezAgentStore {
     addGroupMember: (groupId: string, memberId: string) => void;
     removeGroupMember: (groupId: string, memberId: string) => void;
     getSwarmAgents: () => WaldiezNodeAgentSwarm[];
+    setSwarmInitialAgent: (agentId: string) => void;
+    updateSwarmInitialAgent: (agentId: string) => void;
     getNonSwarmAgents: (
         swarmContainerId: string,
         swarmAgents: WaldiezNodeAgent[],
diff --git a/src/waldiez/models/Stores/IFlowStore.ts b/src/waldiez/models/Stores/IFlowStore.ts
index 177b270..5921172 100644
--- a/src/waldiez/models/Stores/IFlowStore.ts
+++ b/src/waldiez/models/Stores/IFlowStore.ts
@@ -17,7 +17,7 @@ export interface IWaldiezFlowStore {
     onFlowChanged: () => WaldiezFlow;
     onViewportChange: (viewport: { x: number; y: number; zoom: number }, nodeType: WaldiezNodeType) => void;
     saveFlow: () => void;
-    getFlowEdges: () => { used: WaldiezEdge[]; remaining: WaldiezEdge[] };
+    getFlowEdges: (skipSwarm: boolean) => { used: WaldiezEdge[]; remaining: WaldiezEdge[] };
     importFlow: (items: ThingsToImport, flowData: ImportedFlow, typeShown: WaldiezNodeType) => void;
     exportFlow: (hideSecrets: boolean, skipLinks: boolean) => WaldiezFlow;
     updateFlowOrder: (data: { id: string; order: number }[]) => void;
diff --git a/src/waldiez/store/agent.ts b/src/waldiez/store/agent.ts
index 8dad424..953b857 100644
--- a/src/waldiez/store/agent.ts
+++ b/src/waldiez/store/agent.ts
@@ -8,7 +8,13 @@ import {
 } from "@waldiez/models";
 import { IWaldiezAgentStore, WaldiezChat } from "@waldiez/models";
 import { agentMapper, chatMapper } from "@waldiez/models/mappers";
-import { getAgentConnections, getAgentNode, resetEdgeOrders, resetEdgePositions } from "@waldiez/store/utils";
+import {
+    getAgentConnections,
+    getAgentNode,
+    resetEdgeOrders,
+    resetEdgePositions,
+    setSwarmInitialAgent,
+} from "@waldiez/store/utils";
 import { typeOfGet, typeOfSet } from "@waldiez/types";
 import { getId } from "@waldiez/utils";
 
@@ -229,6 +235,40 @@ export class WaldiezAgentStore implements IWaldiezAgentStore {
             node => node.type === "agent" && node.data.agentType === "swarm",
         ) as WaldiezNodeAgentSwarm[];
     };
+    setSwarmInitialAgent = (agentId: string) => {
+        setSwarmInitialAgent(agentId, this.get, this.set);
+    };
+    updateSwarmInitialAgent = (agentId: string) => {
+        setSwarmInitialAgent(agentId, this.get, this.set);
+        // if there are any edges (with source a non-swarm agent)
+        // that connect to a swarm agent, update the edge to connect to the new initial agent
+        // (it should connect to the old one, check this too?)
+        this.set({
+            edges: this.get().edges.map(edge => {
+                const sourceNode = this.get().nodes.find(node => node.id === edge.source);
+                if (sourceNode && sourceNode.data.agentType !== "swarm") {
+                    const targetNode = this.get().nodes.find(node => node.id === edge.target);
+                    if (targetNode && targetNode.data.agentType === "swarm") {
+                        // id={`agent-handle-bottom-target-${id}`}
+                        // id={`agent-handle-top-target-${id}`}...
+                        // if the old handle is in the format above,
+                        // let's update it to the new initial agent's id
+                        let newTargetHandle;
+                        const oldTargetHandle = edge.targetHandle ?? "";
+                        const targetHandleParts = oldTargetHandle.split("-");
+                        if (targetHandleParts.length === 5) {
+                            newTargetHandle = `${targetHandleParts[0]}-${targetHandleParts[1]}-${targetHandleParts[2]}-${
+                                targetHandleParts[3]
+                            }-${agentId}`;
+                        }
+                        return { ...edge, target: agentId, targetHandle: newTargetHandle };
+                    }
+                }
+                return edge;
+            }),
+            updatedAt: new Date().toISOString(),
+        });
+    };
     getNonSwarmAgents: (
         swarmContainerId: string,
         swarmAgents: WaldiezNodeAgent[],
diff --git a/src/waldiez/store/edge.ts b/src/waldiez/store/edge.ts
index 2477c8a..b6d6bc1 100644
--- a/src/waldiez/store/edge.ts
+++ b/src/waldiez/store/edge.ts
@@ -12,8 +12,10 @@ import {
     edgeCommonStyle,
     getNewEdge,
     getNewEdgeConnectionProps,
+    getNewEdgeNodes,
     resetEdgeOrders,
     resetEdgePositions,
+    setSwarmInitialAgent,
     shouldReconnect,
 } from "@waldiez/store/utils";
 import { AGENT_COLORS } from "@waldiez/theme";
@@ -139,16 +141,24 @@ export class WaldiezEdgeStore implements IWaldiezEdgeStore {
         const nodes = this.get().nodes as WaldiezNodeAgent[];
         const edges = this.get().edges;
         const edgesCounter = (chatType: string) => edges.filter(edge => edge.type === chatType).length;
-        const newEdge = getNewEdge(connection, hidden, edgesCounter, nodes, edges);
+        const { source, target, sourceHandle, targetHandle } = connection;
+        const { sourceNode, targetNode } = getNewEdgeNodes(nodes, source, target);
+        if (!sourceNode || !targetNode) {
+            return null;
+        }
+        const newEdge = getNewEdge(hidden, edgesCounter, sourceNode, targetNode, edges);
         if (!newEdge) {
             return null;
         }
         this.set({
-            edges: [...this.get().edges, { ...newEdge }],
+            edges: [...this.get().edges, { ...newEdge, sourceHandle, targetHandle }],
             updatedAt: new Date().toISOString(),
         });
         this.resetEdgePositions();
         const newStoredEdge = this.get().edges.find(edge => edge.id === newEdge.id);
+        if (sourceNode.data.agentType !== "swarm" && targetNode.data.agentType === "swarm") {
+            setSwarmInitialAgent(targetNode.id, this.get, this.set);
+        }
         return (newStoredEdge ?? newEdge) as WaldiezEdge;
     };
     onEdgeDoubleClick = (_event: any, edge: WaldiezEdge) => {
@@ -191,8 +201,7 @@ export class WaldiezEdgeStore implements IWaldiezEdgeStore {
     };
     onReconnect: (oldEdge: Edge, newConnection: Connection) => void = (oldEdge, newConnection) => {
         const nodes = this.get().nodes as WaldiezNodeAgent[];
-        const edges = this.get().edges as WaldiezEdge[];
-        if (!shouldReconnect(newConnection, nodes, edges)) {
+        if (!shouldReconnect(newConnection, nodes)) {
             return;
         }
         const { oldSourceNode, oldTargetNode, newSourceNode, newTargetNode, color } =
@@ -201,8 +210,8 @@ export class WaldiezEdgeStore implements IWaldiezEdgeStore {
             console.error("Not all nodes found");
             return;
         }
-        if (oldSourceNode.id === newSourceNode.id && oldTargetNode.id === newTargetNode.id) {
-            return;
+        if (oldSourceNode.data.agentType !== "swarm" && newTargetNode.data.agentType === "swarm") {
+            setSwarmInitialAgent(newTargetNode.id, this.get, this.set);
         }
         if (!color) {
             return false;
diff --git a/src/waldiez/store/flow.ts b/src/waldiez/store/flow.ts
index 04bf508..325f513 100644
--- a/src/waldiez/store/flow.ts
+++ b/src/waldiez/store/flow.ts
@@ -68,10 +68,10 @@ export class WaldiezFlowStore implements IWaldiezFlowStore {
         }
         return exported;
     };
-    getFlowEdges = () => {
-        const allEdges = this.get().edges.filter(
-            edge => edge.type === "chat" || edge.type === "swarm",
-        ) as WaldiezEdge[];
+    getFlowEdges = (skipSwarm: boolean) => {
+        const allEdges = skipSwarm
+            ? this.get().edges.filter(edge => edge.type === "chat")
+            : this.get().edges.filter(edge => edge.type === "chat" || edge.type === "swarm");
         const usedEdges = [] as WaldiezEdge[];
         const remainingEdges = [] as WaldiezEdge[];
         allEdges.forEach(edge => {
@@ -82,9 +82,9 @@ export class WaldiezFlowStore implements IWaldiezFlowStore {
                 edgeOrder = -1;
             }
             if (edgeOrder >= 0) {
-                usedEdges.push(edge);
+                usedEdges.push(edge as WaldiezEdge);
             } else {
-                remainingEdges.push(edge);
+                remainingEdges.push(edge as WaldiezEdge);
             }
         });
         const sortedEdgesUsed = usedEdges.sort((a, b) => (a.data?.order ?? 0) - (b.data?.order ?? 0));
diff --git a/src/waldiez/store/utils/agent.ts b/src/waldiez/store/utils/agent.ts
index d44e125..0b8de67 100644
--- a/src/waldiez/store/utils/agent.ts
+++ b/src/waldiez/store/utils/agent.ts
@@ -11,6 +11,7 @@ import {
     WaldiezNodeAgentType,
     agentMapper,
 } from "@waldiez/models";
+import { typeOfGet, typeOfSet } from "@waldiez/types";
 
 export const getAgentNode = (
     agentType: WaldiezNodeAgentType,
@@ -103,3 +104,38 @@ const getAgentEdgeConnections = (
     }
     return { sourceNode, targetNode };
 };
+
+export const setSwarmInitialAgent = (agentId: string, get: typeOfGet, set: typeOfSet) => {
+    set({
+        nodes: get().nodes.map(node => {
+            if (node.data.agentType === "swarm") {
+                if (node.id === agentId) {
+                    return {
+                        ...node,
+                        data: {
+                            ...node.data,
+                            isInitial: true,
+                        },
+                    };
+                }
+                return {
+                    ...node,
+                    data: {
+                        ...node.data,
+                        isInitial: false,
+                    },
+                };
+            } else if (node.data.agentType === "swarm_container") {
+                return {
+                    ...node,
+                    data: {
+                        ...node.data,
+                        initialAgent: agentId,
+                    },
+                };
+            }
+            return node;
+        }),
+        updatedAt: new Date().toISOString(),
+    });
+};
diff --git a/src/waldiez/store/utils/edge.ts b/src/waldiez/store/utils/edge.ts
index 7493cc2..cac3107 100644
--- a/src/waldiez/store/utils/edge.ts
+++ b/src/waldiez/store/utils/edge.ts
@@ -22,14 +22,16 @@ export const edgeCommonStyle = (edgeType: WaldiezEdgeType, color: string) => ({
     },
 });
 
-const getNewEdgeNodes = (allNodes: Node[], source: string, target: string) => {
+export const getNewEdgeNodes = (allNodes: Node[], source: string, target: string) => {
     const sourceNode = allNodes.find(node => node.id === source);
     if (!sourceNode) {
-        throw new Error(`Source node with id ${source} not found`);
+        // throw new Error(`Source node with id ${source} not found`);
+        return { sourceNode: null, targetNode: null };
     }
     const targetNode = allNodes.find(node => node.id === target);
     if (!targetNode) {
-        throw new Error(`Target node with id ${target} not found`);
+        // throw new Error(`Target node with id ${target} not found`);
+        return { sourceNode: null, targetNode: null };
     }
     return {
         sourceNode,
@@ -49,7 +51,7 @@ const isSwarmEdge = (sourceNode: Node, targetNode: Node) => {
     if (sourceAgentType === "swarm") {
         return true;
     }
-    if (["user", "rag_user"].includes(sourceAgentType) && targetAgentType === "swarm_container") {
+    if (["user", "rag_user"].includes(sourceAgentType) && targetAgentType === "swarm") {
         return true;
     }
     return false;
@@ -74,16 +76,11 @@ const getNewChatType: (sourceNode: Node, targetNode: Node, hidden: boolean) => W
     return chatType;
 };
 
-const getSwarmEdge = (
-    edges: Edge[],
-    sourceNode: Node,
-    targetNode: Node,
-    sourceHandle: string | null,
-    targetHandle: string | null,
-) => {
-    if (targetNode.data.agentType === "swarm_container") {
-        const existingEdge = edges.find(edge => edge.target === targetNode.id);
-        if (existingEdge) {
+const getSwarmEdge = (edges: Edge[], sourceNode: Node, targetNode: Node) => {
+    // if the source is not swarm, it should not connect to any other node
+    if (sourceNode.data.agentType !== "swarm") {
+        const existing = edges.find(edge => edge.source === sourceNode.id);
+        if (existing) {
             return null;
         }
     }
@@ -102,8 +99,6 @@ const getSwarmEdge = (
     return {
         ...newEdge,
         type: "swarm",
-        sourceHandle,
-        targetHandle,
         animated,
         selected: true,
         ...edgeCommonStyle("swarm", color),
@@ -152,24 +147,22 @@ const getSwarmChatData = (sourceNode: Node, targetNode: Node) => {
 };
 
 export const getNewEdge = (
-    connection: Connection,
     hidden: boolean,
     positionGetter: (chatType: string) => number,
-    nodes: Node[],
+    sourceNode: Node,
+    targetNode: Node,
     edges: Edge[],
 ) => {
-    const { source, target, sourceHandle, targetHandle } = connection;
-    const { sourceNode, targetNode } = getNewEdgeNodes(nodes, source, target);
     if (isSwarmEdge(sourceNode, targetNode)) {
-        return getSwarmEdge(edges, sourceNode, targetNode, sourceHandle, targetHandle);
+        return getSwarmEdge(edges, sourceNode, targetNode);
     }
     if (["swarm", "swarm_container"].includes(targetNode.data.agentType as WaldiezNodeAgentType)) {
         return null;
     }
     const edgeName = getNewEdgeName(sourceNode, targetNode);
     const chatData = new WaldiezChatData();
-    chatData.source = source;
-    chatData.target = target;
+    chatData.source = sourceNode.id;
+    chatData.target = targetNode.id;
     chatData.name = edgeName;
     chatData.order = -1;
     const chatType = getNewChatType(sourceNode, targetNode, hidden);
@@ -185,8 +178,6 @@ export const getNewEdge = (
     return {
         ...newEdge,
         type: chatType,
-        sourceHandle,
-        targetHandle,
         animated: false,
         selected: true,
         ...edgeCommonStyle(chatType, color),
@@ -268,7 +259,7 @@ export const resetEdgePositions = (get: typeOfGet, set: typeOfSet) => {
     });
     updateNestedEdges(get, set);
 };
-export const shouldReconnect = (newConnection: Connection, nodes: Node[], edges: Edge[]): boolean => {
+export const shouldReconnect = (newConnection: Connection, nodes: Node[]): boolean => {
     const newTarget = nodes.find(node => node.id === newConnection.target);
     const newSource = nodes.find(node => node.id === newConnection.source);
     if (!newSource || !newTarget) {
@@ -276,26 +267,19 @@ export const shouldReconnect = (newConnection: Connection, nodes: Node[], edges:
     }
     const isSourceSwarmAgent = newSource.data.agentType === "swarm";
     const isTargetSwarmAgent = newTarget.data.agentType === "swarm";
-    const isTargetSwarmContainer = newTarget.data.agentType === "swarm_container";
     // if not a swarm connection, allow reconnect
-    if (!isSourceSwarmAgent && !isTargetSwarmAgent && !isTargetSwarmContainer) {
+    if (!isSourceSwarmAgent && !isTargetSwarmAgent) {
         return true;
     }
-    if (isSourceSwarmAgent && !isTargetSwarmContainer) {
+    if (isSourceSwarmAgent && ["swarm", "assistant"].includes(newTarget.data.agentType as string)) {
         return true;
     }
-    // if the target is a swarm agent, the source should be a swarm agent too
-    if (isTargetSwarmAgent && !isSourceSwarmAgent) {
-        return false;
-    }
-    // if the target is a swarm container, only allow one edge to it
-    if (isTargetSwarmContainer) {
-        const existingSwarmEdgeToContainer = edges.find(edge => edge.target === newTarget.id);
-        if (existingSwarmEdgeToContainer) {
-            return false;
-        }
-        return true;
+    // if the source is not a swarm agent, but the target is,
+    // only allow user/rag_user to connect to swarm
+    if (!isSourceSwarmAgent && isTargetSwarmAgent) {
+        return ["user", "rag_user"].includes(newSource.data.agentType as string);
     }
+    // if both are swarm agents, allow reconnect
     return isSourceSwarmAgent && isTargetSwarmAgent;
 };
 export const getNewEdgeConnectionProps = (
diff --git a/src/waldiez/styles/overrides.css b/src/waldiez/styles/overrides.css
index 5dc42ab..d650d57 100644
--- a/src/waldiez/styles/overrides.css
+++ b/src/waldiez/styles/overrides.css
@@ -12,9 +12,9 @@
 }
 
 .flow-wrapper .agent-node .react-flow__handle {
-    width: 10px;
-    height: 10px;
-    border-radius: 2px;
+    width: 12px;
+    height: 12px;
+    border-radius: 50%;
 }
 
 .flow-wrapper .agent-node .group-member-handle {
@@ -28,80 +28,6 @@
     border: none;
 }
 
-.flow-wrapper .agent-node .react-flow__handle-left,
-.flow-wrapper .agent-node .react-flow__handle-right,
-.flow-wrapper .swarm-container .react-flow__handle-left,
-.flow-wrapper .swarm-container .react-flow__handle-right {
-    width: 12px;
-    height: 12px;
-    border-radius: 50%;
-}
-
-.flow-wrapper .swarm .react-flow__handle {
-    border-radius: 2px;
-    background-color: var(--swarm-agent-color);
-}
-
-.flow-wrapper .swarm .react-flow__handle.target {
-    width: 20px;
-    height: 9px;
-}
-
-.flow-wrapper .swarm .react-flow__handle.source {
-    width: 12px;
-    height: 12px;
-    border-radius: 0;
-    transform: rotate(45deg);
-}
-
-.flow-wrapper .swarm .react-flow__handle-top.source {
-    margin-left: calc(-20%);
-    margin-top: -6px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-top.target {
-    margin-left: calc(20% - 9px);
-    margin-top: 1px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-bottom.target {
-    margin-left: calc(-20% + 6px);
-    margin-bottom: 2px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-bottom.source {
-    margin-left: calc(20% - 12px);
-    margin-bottom: -6px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-left.target {
-    margin-top: calc(-20% + 6px);
-    height: 20px;
-    width: 9px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-left.source {
-    margin-top: calc(20% - 12px);
-    margin-left: -6px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-right.target {
-    margin-top: calc(20% - 6px);
-    height: 20px;
-    width: 9px;
-    margin-right: 1px;
-}
-
-.flow-wrapper .swarm .react-flow__handle-right.source {
-    margin-top: calc(-20%);
-    margin-right: -5px;
-}
-
-.flow-wrapper .agent-node .react-flow__handle-bottom {
-    margin-bottom: -6px;
-    margin-left: -6px;
-}
-
 .flow-wrapper .agent-node.assistant .react-flow__handle {
     background-color: var(--assistant-agent-color);
 }
diff --git a/src/waldiez/styles/swarm.css b/src/waldiez/styles/swarm.css
index ee7415a..8268bec 100644
--- a/src/waldiez/styles/swarm.css
+++ b/src/waldiez/styles/swarm.css
@@ -38,3 +38,26 @@
     align-items: center;
     margin: 5px;
 }
+
+.flow-wrapper .swarm .react-flow__handle {
+    border-radius: 2px;
+    background-color: var(--swarm-agent-color);
+}
+
+.flow-wrapper .swarm .react-flow__handle.target {
+    width: 20px;
+    height: 10px;
+}
+
+.flow-wrapper .swarm .react-flow__handle-left.target,
+.flow-wrapper .swarm .react-flow__handle-right.target {
+    width: 10px;
+    height: 20px;
+}
+
+.flow-wrapper .swarm .react-flow__handle.source {
+    width: 20px;
+    height: 20px;
+    border-radius: 0;
+    clip-path: polygon(50% 0, 100% 50%, 50% 100%, 0 50%);
+}