feat: Improvements and fixes for extension management (#723)

Co-authored-by: Bradley Axen <[email protected]>
block · Jan 24, 2025 · 7be5b12 · 7be5b12
1 parent 33ce191
commit 7be5b12
Show file tree

Hide file tree

Showing 10 changed files with 311 additions and 356 deletions.
diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx
@@ -1,17 +1,26 @@
 import React, { useEffect, useState } from 'react';
+import { addExtensionFromDeepLink } from './extensions';
+import { useNavigate } from 'react-router-dom';
 import LauncherWindow from './LauncherWindow';
 import ChatWindow from './ChatWindow';
 import ErrorScreen from './components/ErrorScreen';
 import 'react-toastify/dist/ReactToastify.css';
 import { ToastContainer } from 'react-toastify';
 import { ModelProvider } from './components/settings/models/ModelContext';
 import { ActiveKeysProvider } from './components/settings/api_keys/ActiveKeysContext';
-import { loadStoredExtensionConfigs } from './extensions';
 
 export default function App() {
   const [fatalError, setFatalError] = useState<string | null>(null);
   const searchParams = new URLSearchParams(window.location.search);
   const isLauncher = searchParams.get('window') === 'launcher';
+  const navigate = useNavigate();
+
+  useEffect(() => {
+    window.electron.on('add-extension', (_, link) => {
+      window.electron.logInfo(`Adding extension from deep link ${link}`);
+      addExtensionFromDeepLink(link, navigate);
+    });
+  }, [navigate]);
 
   useEffect(() => {
     const handleFatalError = (_: any, errorMessage: string) => {
@@ -21,16 +30,6 @@ export default function App() {
     // Listen for fatal errors from main process
     window.electron.on('fatal-error', handleFatalError);
 
-    // Load stored extension configs when the app starts
-    // delay this by a few seconds
-    setTimeout(() => {
-      window.electron.logInfo('App.tsx: Loading stored extension configs');
-      loadStoredExtensionConfigs().catch((error) => {
-        console.error('Failed to load stored extension configs:', error);
-        window.electron.logInfo('App.tsx: Failed to load stored extension configs ' + error);
-      });
-    }, 5000);
-
     return () => {
       window.electron.off('fatal-error', handleFatalError);
     };

diff --git a/ui/desktop/src/ChatWindow.tsx b/ui/desktop/src/ChatWindow.tsx
@@ -1,8 +1,6 @@
 import React, { useEffect, useRef, useState } from 'react';
 import { Message, useChat } from './ai-sdk-fork/useChat';
 import { getApiUrl, getSecretKey } from './config';
-import { extendGoosedFromUrl } from './extensions';
-import { useNavigate } from 'react-router-dom';
 import BottomMenu from './components/BottomMenu';
 import FlappyGoose from './components/FlappyGoose';
 import GooseMessage from './components/GooseMessage';
@@ -24,6 +22,7 @@ import { useRecentModels } from './components/settings/models/RecentModels';
 import { createSelectedModel } from './components/settings/models/utils';
 import { getDefaultModel } from './components/settings/models/hardcoded_stuff';
 import Splash from './components/Splash';
+import { loadAndAddStoredExtensions } from './extensions';
 
 declare global {
   interface Window {
@@ -337,18 +336,6 @@ export default function ChatWindow() {
     };
   }, []);
 
-  const navigate = useNavigate();
-
-  useEffect(() => {
-    // Listen for goose:// deep links
-    window.electron.on('add-extension', (_, link) => {
-      window.electron.logInfo('Received message for add-extension: ' + link);
-      console.log('Received message for add-extension:', link);
-      extendGoosedFromUrl(link, navigate);
-      window.electron.logInfo('extended called: ' + link);
-    });
-  }, [navigate]);
-
   // Get initial query and history from URL parameters
   const searchParams = new URLSearchParams(window.location.search);
   const initialQuery = searchParams.get('initialQuery');

diff --git a/ui/desktop/src/components/settings/Settings.tsx b/ui/desktop/src/components/settings/Settings.tsx
@@ -3,18 +3,23 @@ import { ScrollArea } from '../ui/scroll-area';
 import { useNavigate, useLocation } from 'react-router-dom';
 import { Plus } from 'lucide-react';
 import { Settings as SettingsType } from './types';
-import { FullExtensionConfig, replaceWithShims, extendGoosed } from '../../extensions';
+import {
+  FullExtensionConfig,
+  addExtension,
+  removeExtension,
+  BUILT_IN_EXTENSIONS,
+} from '../../extensions';
 import { ConfigureExtensionModal } from './extensions/ConfigureExtensionModal';
 import { ManualExtensionModal } from './extensions/ManualExtensionModal';
-import { showToast } from '../ui/toast';
 import BackButton from '../ui/BackButton';
 import { RecentModelsRadio } from './models/RecentModels';
 import { ExtensionItem } from './extensions/ExtensionItem';
-import { getApiUrl, getSecretKey } from '../../config';
 
 const EXTENSIONS_DESCRIPTION =
   'The Model Context Protocol (MCP) is a system that allows AI models to securely connect with local or remote resources using standard server setups. It works like a client-server setup and expands AI capabilities using three main components: Prompts, Resources, and Tools.';
 
+const EXTENSIONS_SITE_LINK = 'https://block.github.io/goose/v1/extensions/';
+
 const DEFAULT_SETTINGS: SettingsType = {
   models: [
     {
@@ -36,42 +41,10 @@ const DEFAULT_SETTINGS: SettingsType = {
       enabled: true,
     },
   ],
-  extensions: [],
+  // @ts-expect-error "we actually do always have all the properties required for builtins, but tsc cannot tell for some reason"
+  extensions: BUILT_IN_EXTENSIONS,
 };
 
-const BUILT_IN_EXTENSIONS = [
-  {
-    id: 'jetbrains',
-    name: 'Jetbrains',
-    type: 'stdio',
-    cmd: 'goosed',
-    args: ['mcp', 'jetbrains'],
-    description: 'Integration with any Jetbrains IDE',
-    enabled: false,
-    env_keys: [],
-  },
-  {
-    id: 'nondeveloper',
-    name: 'Non-Developer assistant',
-    type: 'stdio',
-    cmd: 'goosed',
-    args: ['mcp', 'nondeveloper'],
-    description: "General assisant tools that don't require you to be a developer or engineer.",
-    enabled: false,
-    env_keys: [],
-  },
-  {
-    id: 'memory',
-    name: 'Memory',
-    type: 'stdio',
-    cmd: 'goosed',
-    args: ['mcp', 'memory'],
-    description: 'Teach goose your preferences as you go.',
-    enabled: false,
-    env_keys: [],
-  },
-];
-
 export default function Settings() {
   const navigate = useNavigate();
   const location = useLocation();
@@ -96,6 +69,7 @@ export default function Settings() {
 
   const [extensionBeingConfigured, setExtensionBeingConfigured] =
     useState<FullExtensionConfig | null>(null);
+
   const [isManualModalOpen, setIsManualModalOpen] = useState(false);
 
   // Persist settings changes
@@ -142,80 +116,32 @@ export default function Settings() {
       ),
     }));
 
-    try {
-      const endpoint = newEnabled ? '/extensions/add' : '/extensions/remove';
-
-      // Full config for adding - only "name" as a string for removing
-      const body = newEnabled
-        ? {
-            type: extension.type,
-            ...(extension.type === 'stdio' && {
-              cmd: await replaceWithShims(extension.cmd),
-              args: extension.args || [],
-            }),
-            ...(extension.type === 'sse' && {
-              uri: extension.uri,
-            }),
-            ...(extension.type === 'builtin' && {
-              name: extension.name,
-            }),
-            env_keys: extension.env_keys,
-          }
-        : extension.name;
-
-      const response = await fetch(getApiUrl(endpoint), {
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          'X-Secret-Key': getSecretKey(),
-        },
-        body: JSON.stringify(body),
-      });
+    let response: Response;
 
-      if (!response.ok) {
-        throw new Error(`Failed to ${newEnabled ? 'enable' : 'disable'} extension`);
-      }
+    if (newEnabled) {
+      response = await addExtension(extension);
+    } else {
+      response = await removeExtension(extension.name);
+    }
 
-      showToast(`Successfully ${newEnabled ? 'enabled' : 'disabled'} extension`, 'success');
-    } catch (error) {
+    if (!response.ok) {
       setSettings(originalSettings);
-      showToast(`Error ${newEnabled ? 'enabling' : 'disabling'} extension`, 'error');
-      console.error('Error toggling extension:', error);
     }
   };
 
   const handleExtensionRemove = async () => {
-    if (!extensionBeingConfigured) return;
-
-    try {
-      // First disable the extension if it's enabled
-      if (extensionBeingConfigured.enabled) {
-        const response = await fetch(getApiUrl('/extensions/remove'), {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/json',
-            'X-Secret-Key': getSecretKey(),
-          },
-          body: JSON.stringify(extensionBeingConfigured.name),
-        });
-
-        if (!response.ok) {
-          throw new Error('Failed to remove extension from backend');
-        }
-      }
+    if (!extensionBeingConfigured || !extensionBeingConfigured.enabled) return;
 
-      // Then remove it from the local settings
+    const response = await removeExtension(extensionBeingConfigured.name);
+
+    if (response.ok) {
+      // Remove from localstorage
       setSettings((prev) => ({
         ...prev,
         extensions: prev.extensions.filter((ext) => ext.id !== extensionBeingConfigured.id),
       }));
-
-      showToast(`Successfully removed ${extensionBeingConfigured.name} extension`, 'success');
       setExtensionBeingConfigured(null);
       navigate('/settings', { replace: true });
-    } catch (error) {
-      console.error('Error removing extension:', error);
-      showToast('Failed to remove extension', 'error');
     }
   };
 
@@ -297,11 +223,7 @@ export default function Settings() {
                       </button>{' '}
                       |
                       <button
-                        onClick={() =>
-                          window.electron.openInChrome(
-                            'https://silver-disco-nvm6v4e.pages.github.io/'
-                          )
-                        }
+                        onClick={() => window.electron.openInChrome(EXTENSIONS_SITE_LINK)}
                         className="text-indigo-500 hover:text-indigo-600 font-medium"
                       >
                         Browse Extensions
@@ -346,37 +268,16 @@ export default function Settings() {
         isOpen={isManualModalOpen}
         onClose={() => setIsManualModalOpen(false)}
         onSubmit={async (extension) => {
-          // Create config for extendGoosed
-          const config = {
-            type: extension.type,
-            ...(extension.type === 'stdio' && {
-              cmd: await replaceWithShims(extension.cmd),
-              args: extension.args || [],
-            }),
-            ...(extension.type === 'sse' && {
-              uri: extension.uri,
-            }),
-            ...(extension.type === 'builtin' && {
-              name: extension.name,
-            }),
-            env_keys: extension.env_keys,
-          };
-
-          try {
-            const success = await extendGoosed(config);
-            if (success) {
-              setSettings((prev) => ({
-                ...prev,
-                extensions: [...prev.extensions, extension],
-              }));
-              setIsManualModalOpen(false);
-              showToast('Extension added successfully', 'success');
-            } else {
-              throw new Error('Failed to add extension');
-            }
-          } catch (error) {
-            console.error('Error adding extension:', error);
-            showToast('Error adding extension', 'error');
+          const response = await addExtension(extension);
+
+          if (response.ok) {
+            setSettings((prev) => ({
+              ...prev,
+              extensions: [...prev.extensions, extension],
+            }));
+            setIsManualModalOpen(false);
+          } else {
+            // TODO - Anything for the UI state beyond validation?
           }
         }}
       />

diff --git a/ui/desktop/src/components/settings/extensions/ConfigureExtensionModal.tsx b/ui/desktop/src/components/settings/extensions/ConfigureExtensionModal.tsx
@@ -2,9 +2,10 @@ import React from 'react';
 import { Card } from '../../ui/card';
 import { Button } from '../../ui/button';
 import { Input } from '../../ui/input';
-import { FullExtensionConfig, replaceWithShims } from '../../../extensions';
+import { FullExtensionConfig } from '../../../extensions';
 import { getApiUrl, getSecretKey } from '../../../config';
-import { showToast } from '../../ui/toast';
+import { addExtension } from '../../../extensions';
+import { toast } from 'react-toastify';
 
 interface ConfigureExtensionModalProps {
   isOpen: boolean;
@@ -62,41 +63,18 @@ export function ConfigureExtensionModal({
         }
       }
 
-      // Then add the system configuration
-      const extensionConfig = {
-        type: extension.type,
-        ...(extension.type === 'stdio' && {
-          cmd: await replaceWithShims(extension.cmd),
-          args: extension.args || [],
-        }),
-        ...(extension.type === 'sse' && {
-          uri: extension.uri,
-        }),
-        ...(extension.type === 'builtin' && {
-          name: extension.name,
-        }),
-        env_keys: extension.env_keys,
-      };
-
-      const response = await fetch(getApiUrl('/extensions/add'), {
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          'X-Secret-Key': getSecretKey(),
-        },
-        body: JSON.stringify(extensionConfig),
-      });
+      const response = await addExtension(extension);
 
       if (!response.ok) {
         throw new Error('Failed to add system configuration');
       }
 
-      showToast(`Successfully configured the ${extension.name} extension`, 'success');
+      toast.success(`Successfully configured the ${extension.name} extension`);
       onSubmit();
       onClose();
     } catch (error) {
       console.error('Error configuring extension:', error);
-      showToast('Failed to configure extension', 'error');
+      toast.error('Failed to configure extension');
     } finally {
       setIsSubmitting(false);
     }