ui: start extensions on add (#1714)

Co-authored-by: Ben Walding <bwalding@cloudbees.com>

ui/desktop/src/agent/UpdateAgent.tsx

@@ -0,0 +1,182 @@
+import { useConfig, FixedExtensionEntry } from '../components/ConfigContext';
+import { getApiUrl, getSecretKey } from '../config';
+import { ExtensionConfig } from '../api';
+import { toast } from 'react-toastify';
+import React, { useState } from 'react';
+import { initializeAgent as startAgent, replaceWithShims } from './utils';
+
+// extensionUpdate = an extension was newly added or updated so we should attempt to add it
+
+export const useAgent = () => {
+  const { getExtensions, read } = useConfig();
+  const [isUpdating, setIsUpdating] = useState(false);
+
+  // whenever we change the model, we must call this
+  const initializeAgent = async (provider: string, model: string) => {
+    try {
+      console.log('Initializing agent with provider', provider, 'model', model);
+
+      const response = await startAgent(model, provider);
+
+      if (!response.ok) {
+        throw new Error(`Failed to initialize agent: ${response.statusText}`);
+      }
+
+      return true;
+    } catch (error) {
+      console.error('Failed to initialize agent:', error);
+      toast.error(
+        `Failed to initialize agent: ${error instanceof Error ? error.message : 'Unknown error'}`
+      );
+      return false;
+    }
+  };
+
+  const updateAgent = async (extensionUpdate?: ExtensionConfig) => {
+    setIsUpdating(true);
+
+    try {
+      // need to initialize agent first (i dont get why but if we dont do this, we get a 428)
+      // note: we must write the value for GOOSE_MODEL and GOOSE_PROVIDER in the config before updating agent
+      const goose_model = (await read('GOOSE_MODEL', false)) as string;
+      const goose_provider = (await read('GOOSE_PROVIDER', false)) as string;
+
+      console.log(
+        `Starting agent with GOOSE_MODEL=${goose_model} and GOOSE_PROVIDER=${goose_provider}`
+      );
+
+      // Initialize the agent if it's a model change
+      if (goose_model && goose_provider) {
+        const success = await initializeAgent(goose_provider, goose_model);
+        if (!success) {
+          console.error('Failed to initialize agent during model change');
+          return false;
+        }
+      }
+
+      if (extensionUpdate) {
+        await addExtensionToAgent(extensionUpdate);
+      }
+
+      return true;
+    } catch (error) {
+      console.error('Error updating agent:', error);
+      return false;
+    } finally {
+      setIsUpdating(false);
+    }
+  };
+
+  // TODO: set 'enabled' to false if we fail to start / add the extension
+  // only for non-builtins
+
+  // TODO: try to add some descriptive error messages for common failure modes
+  const addExtensionToAgent = async (
+    extension: ExtensionConfig,
+    silent: boolean = false
+  ): Promise<Response> => {
+    if (extension.type == 'stdio') {
+      console.log('extension command', extension.cmd);
+      extension.cmd = await replaceWithShims(extension.cmd);
+      console.log('next ext command', extension.cmd);
+    }
+
+    try {
+      let toastId;
+      if (!silent) {
+        toastId = toast.loading(`Adding ${extension.name} extension...`, {
+          position: 'top-center',
+        });
+        toast.info('Press the escape key to continue using goose while extension loads');
+      }
+
+      const response = await fetch(getApiUrl('/extensions/add'), {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'X-Secret-Key': getSecretKey(),
+        },
+        body: JSON.stringify(extension),
+      });
+
+      // Handle non-OK responses
+      if (!response.ok) {
+        const errorMsg = `Server returned ${response.status}: ${response.statusText}`;
+        console.error(errorMsg);
+
+        // Special handling for 428 Precondition Required (agent not initialized)
+        if (response.status === 428) {
+          if (!silent) {
+            if (toastId) toast.dismiss(toastId);
+            toast.error('Agent is not initialized. Please initialize the agent first.');
+          }
+          return response;
+        }
+
+        if (!silent) {
+          if (toastId) toast.dismiss(toastId);
+          toast.error(`Failed to add ${extension.name} extension: ${errorMsg}`);
+        }
+        return response;
+      }
+
+      // Parse response JSON safely
+      let data;
+      try {
+        const text = await response.text();
+        data = text ? JSON.parse(text) : { error: false };
+      } catch (error) {
+        console.warn('Could not parse response as JSON, assuming success', error);
+        data = { error: false };
+      }
+
+      console.log('Response data:', data);
+
+      if (!data.error) {
+        if (!silent) {
+          if (toastId) toast.dismiss(toastId);
+          toast.success(`Successfully enabled ${extension.name} extension`);
+        }
+        return response;
+      }
+
+      console.log('Error trying to send a request to the extensions endpoint');
+      const errorMessage = `Error adding ${extension.name} extension${data.message ? `. ${data.message}` : ''}`;
+      const ErrorMsg = ({ closeToast }: { closeToast?: () => void }) => (
+        <div className="flex flex-col gap-1">
+          <div>Error adding {extension.name} extension</div>
+          <div>
+            <button
+              className="text-sm rounded px-2 py-1 bg-gray-400 hover:bg-gray-300 text-white cursor-pointer"
+              onClick={() => {
+                navigator.clipboard.writeText(data.message || 'Unknown error');
+                closeToast?.();
+              }}
+            >
+              Copy error message
+            </button>
+          </div>
+        </div>
+      );
+
+      console.error(errorMessage);
+      if (toastId) toast.dismiss(toastId);
+      toast(ErrorMsg, { type: 'error', autoClose: false });
+
+      return response;
+    } catch (error) {
+      console.log('Got some other error');
+      const errorMessage = `Failed to add ${extension.name} extension: ${error instanceof Error ? error.message : 'Unknown error'}`;
+      console.error(errorMessage);
+      toast.error(errorMessage, { autoClose: false });
+      throw error;
+    }
+  };
+
+  return {
+    updateAgent,
+    addExtensionToAgent,
+    initializeAgent,
+    isUpdating,
+  };
+};

ui/desktop/src/agent/utils.tsx

@@ -0,0 +1,32 @@
+import { getApiUrl, getSecretKey } from '../config';
+
+export async function initializeAgent(model: string, provider: string) {
+  console.log('fetching...', provider, model);
+  const response = await fetch(getApiUrl('/agent'), {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'X-Secret-Key': getSecretKey(),
+    },
+    body: JSON.stringify({
+      provider: provider.toLowerCase().replace(/ /g, '_'),
+      model: model,
+    }),
+  });
+  return response;
+}
+
+export async function replaceWithShims(cmd: string) {
+  const binaryPathMap: Record<string, string> = {
+    goosed: await window.electron.getBinaryPath('goosed'),
+    npx: await window.electron.getBinaryPath('npx'),
+    uvx: await window.electron.getBinaryPath('uvx'),
+  };
+
+  if (binaryPathMap[cmd]) {
+    console.log('--------> Replacing command with shim ------>', cmd, binaryPathMap[cmd]);
+    cmd = binaryPathMap[cmd];
+  }
+
+  return cmd;
+}

ui/desktop/src/components/settings_v2/extensions/ExtensionsSection.tsx

@@ -11,6 +11,7 @@ import {
   extensionToFormData,
   getDefaultFormData,
 } from './utils';
+import { useAgent } from '../../../agent/UpdateAgent';
 
 export default function ExtensionsSection() {
   const { toggleExtension, getExtensions, addExtension, removeExtension } = useConfig();
@@ -20,6 +21,7 @@ export default function ExtensionsSection() {
   const [selectedExtension, setSelectedExtension] = useState<FixedExtensionEntry | null>(null);
   const [isModalOpen, setIsModalOpen] = useState(false);
   const [isAddModalOpen, setIsAddModalOpen] = useState(false);
+  const { updateAgent, addExtensionToAgent } = useAgent();
 
   const fetchExtensions = async () => {
     setLoading(true);
@@ -60,8 +62,10 @@ export default function ExtensionsSection() {
 
     try {
       await addExtension(formData.name, extensionConfig, formData.enabled);
+      console.log('attempting to add extension');
+      await updateAgent(extensionConfig);
       handleModalClose();
-      fetchExtensions(); // Refresh the list after adding
+      await fetchExtensions(); // Refresh the list after adding
     } catch (error) {
       console.error('Failed to add extension:', error);
     }

ui/desktop/src/components/settings_v2/providers/ProviderSettingsPage.tsx

@@ -4,6 +4,7 @@ import BackButton from '../../ui/BackButton';
 import ProviderGrid from './ProviderGrid';
 import { useConfig } from '../../ConfigContext';
 import { ProviderDetails } from '../../../api/types.gen';
+import { useAgent } from '../../../agent/UpdateAgent';
 
 interface ProviderSettingsProps {
   onClose: () => void;
@@ -12,6 +13,7 @@ interface ProviderSettingsProps {
 
 export default function ProviderSettings({ onClose, isOnboarding }: ProviderSettingsProps) {
   const { getProviders, upsert } = useConfig();
+  const { initializeAgent } = useAgent();
   const [loading, setLoading] = useState(true);
   const [providers, setProviders] = useState<ProviderDetails[]>([]);
   const initialLoadDone = useRef(false);
@@ -50,24 +52,30 @@ export default function ProviderSettings({ onClose, isOnboarding }: ProviderSett
 
   // Handler for when a provider is launched if this component is used as part of onboarding page
   const handleProviderLaunch = useCallback(
-    (provider: ProviderDetails) => {
+    async (provider: ProviderDetails) => {
+      const provider_name = provider.name;
+      const model = provider.metadata.default_model;
+
       console.log(`Launching with provider: ${provider.name}`);
       try {
+        // update the config
         // set GOOSE_PROVIDER in the config file
-        // @lily-de: leaving as test for now to avoid messing with my config directly
-        upsert('GOOSE_PROVIDER_TEST', provider.name, false).then((_) =>
-          console.log('Setting GOOSE_PROVIDER to', provider.name)
+        upsert('GOOSE_PROVIDER', provider_name, false).then((_) =>
+          console.log('Setting GOOSE_PROVIDER to', provider_name)
         );
         // set GOOSE_MODEL in the config file
-        upsert('GOOSE_MODEL_TEST', provider.metadata.default_model, false).then((_) =>
-          console.log('Setting GOOSE_MODEL to', provider.metadata.default_model)
+        upsert('GOOSE_MODEL', model, false).then((_) =>
+          console.log('Setting GOOSE_MODEL to', model)
         );
+
+        // initialize agent
+        await initializeAgent(provider_name, model);
       } catch (error) {
-        console.error(`Failed to initialize with provider ${provider.name}:`, error);
+        console.error(`Failed to initialize with provider ${provider_name}:`, error);
       }
       onClose();
     },
-    [onClose, upsert]
+    [initializeAgent, onClose, upsert]
   );
 
   return (