From b22d173425e51e9de91d2826cfd51682ee24b4d7 Mon Sep 17 00:00:00 2001
From: Lily Delalande <119957291+lily-de@users.noreply.github.com>
Date: Fri, 28 Mar 2025 17:10:18 -0400
Subject: [PATCH] ui: toggle back off after failure (#1911)

---
 .../extensions/ExtensionsSection.tsx          | 22 ++++++---
 .../subcomponents/ExtensionItem.tsx           | 45 ++++++++++++++++---
 .../subcomponents/ExtensionList.tsx           |  3 +-
 3 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/ui/desktop/src/components/settings_v2/extensions/ExtensionsSection.tsx b/ui/desktop/src/components/settings_v2/extensions/ExtensionsSection.tsx
index 82fd5934..534fdd69 100644
--- a/ui/desktop/src/components/settings_v2/extensions/ExtensionsSection.tsx
+++ b/ui/desktop/src/components/settings_v2/extensions/ExtensionsSection.tsx
@@ -49,14 +49,22 @@ export default function ExtensionsSection() {
     const toggleDirection = extension.enabled ? 'toggleOff' : 'toggleOn';
     const extensionConfig = extractExtensionConfig(extension);
 
-    await toggleExtension({
-      toggle: toggleDirection,
-      extensionConfig: extensionConfig,
-      addToConfig: addExtension,
-      toastOptions: { silent: false },
-    });
+    try {
+      await toggleExtension({
+        toggle: toggleDirection,
+        extensionConfig: extensionConfig,
+        addToConfig: addExtension,
+        toastOptions: { silent: false },
+      });
 
-    await fetchExtensions(); // Refresh the list after toggling
+      await fetchExtensions(); // Refresh the list after successful toggle
+      return true; // Indicate success
+    } catch (error) {
+      console.error('Toggle extension failed:', error);
+      // Don't refresh the extension list on failure - this allows our visual state rollback to work
+      // The actual state in the config hasn't changed anyway
+      throw error; // Re-throw to let the ExtensionItem component know it failed
+    }
   };
 
   const handleConfigureClick = (extension: FixedExtensionEntry) => {
diff --git a/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionItem.tsx b/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionItem.tsx
index b5e286e5..1ebc0c44 100644
--- a/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionItem.tsx
+++ b/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionItem.tsx
@@ -1,5 +1,4 @@
-// ExtensionItem.tsx
-import React from 'react';
+import React, { useState, useEffect } from 'react';
 import { Switch } from '../../../ui/switch';
 import { Gear } from '../../../icons/Gear';
 import { FixedExtensionEntry } from '../../../ConfigContext';
@@ -7,11 +6,46 @@ import { getSubtitle, getFriendlyTitle } from './ExtensionList';
 
 interface ExtensionItemProps {
   extension: FixedExtensionEntry;
-  onToggle: (extension: FixedExtensionEntry) => void;
+  onToggle: (extension: FixedExtensionEntry) => Promise<boolean | void>;
   onConfigure: (extension: FixedExtensionEntry) => void;
 }
 
 export default function ExtensionItem({ extension, onToggle, onConfigure }: ExtensionItemProps) {
+  // Add local state to track the visual toggle state
+  const [visuallyEnabled, setVisuallyEnabled] = useState(extension.enabled);
+  // Track if we're in the process of toggling
+  const [isToggling, setIsToggling] = useState(false);
+
+  const handleToggle = async (ext: FixedExtensionEntry) => {
+    // Prevent multiple toggles while one is in progress
+    if (isToggling) return;
+
+    setIsToggling(true);
+
+    // Immediately update visual state
+    const newState = !ext.enabled;
+    setVisuallyEnabled(newState);
+
+    try {
+      // Call the actual toggle function that performs the async operation
+      await onToggle(ext);
+      // Success case is handled by the useEffect below when extension.enabled changes
+    } catch (error) {
+      // If there was an error, revert the visual state
+      console.log('Toggle failed, reverting visual state');
+      setVisuallyEnabled(!newState);
+    } finally {
+      setIsToggling(false);
+    }
+  };
+
+  // Update visual state when the actual extension state changes
+  useEffect(() => {
+    if (!isToggling) {
+      setVisuallyEnabled(extension.enabled);
+    }
+  }, [extension.enabled, isToggling]);
+
   const renderFormattedSubtitle = () => {
     const subtitle = getSubtitle(extension);
     return subtitle.split('\n').map((part, index) => (
@@ -21,6 +55,7 @@ export default function ExtensionItem({ extension, onToggle, onConfigure }: Exte
       </React.Fragment>
     ));
   };
+
   return (
     <div className="rounded-lg border border-borderSubtle p-4 mb-2">
       <div className="flex items-center justify-between mb-2">
@@ -36,8 +71,8 @@ export default function ExtensionItem({ extension, onToggle, onConfigure }: Exte
             </button>
           )}
           <Switch
-            checked={extension.enabled}
-            onCheckedChange={() => onToggle(extension)}
+            checked={(isToggling && visuallyEnabled) || extension.enabled}
+            onCheckedChange={() => handleToggle(extension)}
             variant="mono"
           />
         </div>
diff --git a/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionList.tsx b/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionList.tsx
index a36c7e60..c92d8bcf 100644
--- a/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionList.tsx
+++ b/ui/desktop/src/components/settings_v2/extensions/subcomponents/ExtensionList.tsx
@@ -7,7 +7,7 @@ import { combineCmdAndArgs } from '../utils';
 
 interface ExtensionListProps {
   extensions: FixedExtensionEntry[];
-  onToggle: (extension: FixedExtensionEntry) => void;
+  onToggle: (extension: FixedExtensionEntry) => Promise<boolean | void>;
   onConfigure: (extension: FixedExtensionEntry) => void;
 }
 
@@ -25,7 +25,6 @@ export default function ExtensionList({ extensions, onToggle, onConfigure }: Ext
     </div>
   );
 }
-
 // Helper functions
 // Helper function to get a friendly title from extension name
 export function getFriendlyTitle(extension: FixedExtensionEntry): string {