diff --git a/src/components/Flow.tsx b/src/components/Flow.tsx
index f7d1889..90699cd 100644
--- a/src/components/Flow.tsx
+++ b/src/components/Flow.tsx
@@ -641,6 +641,15 @@ export default function App() {
     [setEdges],
   )
 
+  const handleEdgeUnselect = (edgeId: string) => {
+    setEdges((eds) =>
+      eds.map((edge) => ({
+        ...edge,
+        selected: edge.id === edgeId ? false : edge.selected,
+      })),
+    )
+  }
+
   const flowNodes =
     !initialOnboardingComplete &&
     currentOnboardingStep < onboardingSteps.length &&
@@ -653,7 +662,13 @@ export default function App() {
     currentOnboardingStep < onboardingSteps.length &&
     onboardingSteps[currentOnboardingStep].edges
       ? onboardingSteps[currentOnboardingStep].edges
-      : edges
+      : edges.map((edge) => ({
+          ...edge,
+          data: {
+            ...edge.data,
+            onEdgeUnselect: handleEdgeUnselect,
+          },
+        }))
 
   function generateSpec(edges: any, currentLanguage: 'python' | 'typescript' = language): string {
     // Step 1: Separate normal edges and animated edges
diff --git a/src/components/edges/SelfConnectingEdge.tsx b/src/components/edges/SelfConnectingEdge.tsx
index b99b6a7..820348e 100644
--- a/src/components/edges/SelfConnectingEdge.tsx
+++ b/src/components/edges/SelfConnectingEdge.tsx
@@ -8,6 +8,7 @@ interface SelfConnectingEdgeProps extends EdgeProps {
   data?: {
     onLabelClick: (id: string) => void
     updateEdgeLabel: (id: string, newLabel: string) => void
+    onEdgeUnselect?: (edgeId: string) => void
   }
 }
 
@@ -15,10 +16,12 @@ const ColorPicker = ({
   color,
   onChange,
   onClose,
+  edgeId,
 }: {
   color: string
   onChange: (e: React.ChangeEvent<HTMLInputElement>) => void
-  onClose: () => void
+  onClose: (edgeId: string) => void
+  edgeId: string
 }) => {
   return createPortal(
     <div className='fixed bottom-5 left-5 z-50' style={{ width: '280px' }}>
@@ -28,7 +31,7 @@ const ColorPicker = ({
           <button
             onClick={(e) => {
               e.stopPropagation()
-              onClose()
+              onClose(edgeId)
             }}
             className='text-sm bg-slate-800 hover:bg-slate-900 text-slate-100 py-1 px-2 rounded-md'
           >
@@ -237,7 +240,15 @@ export default function SelfConnectingEdge(props: SelfConnectingEdgeProps) {
             }}
           />
           {isColorPickerActive && props.selected && (
-            <ColorPicker color={edgeColor} onChange={handleColorChange} onClose={() => setActiveEdgeId(null)} />
+            <ColorPicker
+              color={edgeColor}
+              onChange={handleColorChange}
+              onClose={(edgeId) => {
+                setActiveEdgeId(null)
+                props.data?.onEdgeUnselect?.(edgeId)
+              }}
+              edgeId={id}
+            />
           )}
           {animated &&
             (editingEdgeId === id ? (
@@ -345,7 +356,15 @@ export default function SelfConnectingEdge(props: SelfConnectingEdgeProps) {
           }}
         />
         {isColorPickerActive && (
-          <ColorPicker color={edgeColor} onChange={handleColorChange} onClose={() => setActiveEdgeId(null)} />
+          <ColorPicker
+            color={edgeColor}
+            onChange={handleColorChange}
+            onClose={(edgeId) => {
+              setActiveEdgeId(null)
+              props.data?.onEdgeUnselect?.(edgeId)
+            }}
+            edgeId={id}
+          />
         )}
         {animated &&
           (editingEdgeId === id ? (