feat: set custom models for lead/worker (#4598)

Signed-off-by: Kamal Fariz Mahyuddin <kamal@envoy.com>

ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.test.tsx

@@ -0,0 +1,154 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { render, screen, waitFor, fireEvent } from '@testing-library/react';
+import { LeadWorkerSettings } from './LeadWorkerSettings';
+
+// Mock predefined models utils to force provider-based options (no predefined list)
+vi.mock('../predefinedModelsUtils', () => ({
+  shouldShowPredefinedModels: () => false,
+  getPredefinedModelsFromEnv: () => [],
+}));
+
+// Mocks for useConfig
+const mockRead = vi.fn();
+const mockUpsert = vi.fn();
+const mockRemove = vi.fn();
+const mockGetProviders = vi.fn();
+
+vi.mock('../../../ConfigContext', () => ({
+  useConfig: () => ({
+    read: mockRead,
+    upsert: mockUpsert,
+    remove: mockRemove,
+    getProviders: mockGetProviders,
+  }),
+}));
+
+// Minimal mock for useModelAndProvider
+vi.mock('../../../ModelAndProviderContext', () => ({
+  useModelAndProvider: () => ({
+    currentModel: null,
+  }),
+}));
+
+describe('LeadWorkerSettings', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  const setupHappyPathMocks = () => {
+    // reads
+    mockRead.mockImplementation(async (key: string) => {
+      switch (key) {
+        case 'GOOSE_LEAD_MODEL':
+          return 'my-custom-lead';
+        case 'GOOSE_LEAD_PROVIDER':
+          return 'anthropic';
+        case 'GOOSE_LEAD_TURNS':
+          return 3;
+        case 'GOOSE_LEAD_FAILURE_THRESHOLD':
+          return 2;
+        case 'GOOSE_LEAD_FALLBACK_TURNS':
+          return 2;
+        case 'GOOSE_MODEL':
+          return 'my-custom-worker';
+        case 'GOOSE_PROVIDER':
+          return 'openai';
+        default:
+          return null;
+      }
+    });
+
+    // providers (options do NOT include the custom models above)
+    mockGetProviders.mockResolvedValue([
+      {
+        is_configured: true,
+        name: 'openai',
+        metadata: {
+          display_name: 'OpenAI',
+          known_models: [{ name: 'gpt-4o' }, { name: 'gpt-4o-mini' }],
+        },
+      },
+      {
+        is_configured: true,
+        name: 'anthropic',
+        metadata: {
+          display_name: 'Anthropic',
+          known_models: [{ name: 'claude-3-5-sonnet' }],
+        },
+      },
+    ]);
+
+    // writers
+    mockUpsert.mockResolvedValue(undefined);
+    mockRemove.mockResolvedValue(undefined);
+  };
+
+  it('shows custom inputs for lead/worker when current models are unknown and saves them', async () => {
+    setupHappyPathMocks();
+
+    const onClose = vi.fn();
+    render(<LeadWorkerSettings isOpen={true} onClose={onClose} />);
+
+    // Wait for modal content (not loading)
+    await waitFor(() => {
+      expect(screen.getByText('Lead/Worker Mode')).toBeInTheDocument();
+    });
+
+    // Labels should be present with back-to-list controls
+    await waitFor(() => {
+      expect(screen.getByText('Lead Model')).toBeInTheDocument();
+      expect(screen.getByText('Worker Model')).toBeInTheDocument();
+      // Back to model list appears for each section when in custom mode
+      const backLinks = screen.getAllByText('Back to model list');
+      expect(backLinks.length).toBeGreaterThanOrEqual(2);
+    });
+
+    const inputs = screen.getAllByPlaceholderText('Type model name here') as HTMLInputElement[];
+    expect(inputs.length).toBe(2);
+    const [leadInput, workerInput] = inputs;
+    expect(leadInput.value).toBe('my-custom-lead');
+    expect(workerInput.value).toBe('my-custom-worker');
+
+    // Save settings
+    const saveBtn = screen.getByRole('button', { name: 'Save Settings' });
+    expect(saveBtn).toBeEnabled();
+    fireEvent.click(saveBtn);
+
+    // Assert upserts for models (providers are optional but present in this setup)
+    await waitFor(() => {
+      expect(mockUpsert).toHaveBeenCalledWith('GOOSE_LEAD_MODEL', 'my-custom-lead', false);
+      expect(mockUpsert).toHaveBeenCalledWith('GOOSE_MODEL', 'my-custom-worker', false);
+      expect(mockUpsert).toHaveBeenCalledWith('GOOSE_LEAD_PROVIDER', 'anthropic', false);
+      expect(mockUpsert).toHaveBeenCalledWith('GOOSE_PROVIDER', 'openai', false);
+    });
+  });
+
+  it('disables lead/worker and removes config when toggled off', async () => {
+    setupHappyPathMocks();
+
+    const onClose = vi.fn();
+    render(<LeadWorkerSettings isOpen={true} onClose={onClose} />);
+
+    await waitFor(() => {
+      expect(screen.getByText('Lead/Worker Mode')).toBeInTheDocument();
+    });
+
+    // Toggle off
+    const checkbox = screen.getByLabelText('Enable lead/worker mode') as HTMLInputElement;
+    expect(checkbox.checked).toBe(true);
+    fireEvent.click(checkbox);
+    expect(checkbox.checked).toBe(false);
+
+    const saveBtn = screen.getByRole('button', { name: 'Save Settings' });
+    expect(saveBtn).toBeEnabled();
+    fireEvent.click(saveBtn);
+
+    await waitFor(() => {
+      expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_MODEL', false);
+      expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_PROVIDER', false);
+      expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_TURNS', false);
+      expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_FAILURE_THRESHOLD', false);
+      expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_FALLBACK_TURNS', false);
+    });
+  });
+});

ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx

@@ -19,6 +19,9 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
   const [workerModel, setWorkerModel] = useState<string>('');
   const [leadProvider, setLeadProvider] = useState<string>('');
   const [workerProvider, setWorkerProvider] = useState<string>('');
+  // Minimal custom model mode toggles
+  const [isLeadCustomModel, setIsLeadCustomModel] = useState<boolean>(false);
+  const [isWorkerCustomModel, setIsWorkerCustomModel] = useState<boolean>(false);
   const [leadTurns, setLeadTurns] = useState<number>(3);
   const [failureThreshold, setFailureThreshold] = useState<number>(2);
   const [fallbackTurns, setFallbackTurns] = useState<number>(2);
@@ -113,6 +116,9 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
           });
         }
 
+        // Append a simple "custom" option to enable free-text entry
+        options.push({ value: '__custom__', label: 'Use custom model…', provider: '' });
+
         setModelOptions(options);
       } catch (error) {
         console.error('Error loading configuration:', error);
@@ -124,6 +130,18 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
     loadConfig();
   }, [read, getProviders, currentModel, isOpen]);
 
+  // If current models are not in the list (e.g., previously set to custom), switch to custom mode
+  useEffect(() => {
+    if (!isLoading) {
+      if (leadModel && !modelOptions.find((opt) => opt.value === leadModel)) {
+        setIsLeadCustomModel(true);
+      }
+      if (workerModel && !modelOptions.find((opt) => opt.value === workerModel)) {
+        setIsWorkerCustomModel(true);
+      }
+    }
+  }, [isLoading, modelOptions, leadModel, workerModel]);
+
   const handleSave = async () => {
     try {
       if (isEnabled && leadModel && workerModel) {
@@ -194,46 +212,104 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
 
           <div className="space-y-4">
             <div className="space-y-2">
-              <label className={`text-sm ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
-                Lead Model
-              </label>
-              <Select
-                options={modelOptions}
-                value={modelOptions.find((opt) => opt.value === leadModel) || null}
-                onChange={(newValue: unknown) => {
-                  const option = newValue as { value: string; provider: string } | null;
-                  if (option) {
-                    setLeadModel(option.value);
-                    setLeadProvider(option.provider);
+              <div className="flex items-center justify-between">
+                <label className={`text-sm ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
+                  Lead Model
+                </label>
+                {isLeadCustomModel && (
+                  <button
+                    onClick={() => setIsLeadCustomModel(false)}
+                    className={`text-xs ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'} hover:underline`}
+                    type="button"
+                  >
+                    Back to model list
+                  </button>
+                )}
+              </div>
+              {!isLeadCustomModel ? (
+                <Select
+                  options={modelOptions}
+                  value={
+                    leadModel ? modelOptions.find((opt) => opt.value === leadModel) || null : null
                   }
-                }}
-                placeholder="Select lead model..."
-                isDisabled={!isEnabled}
-                className={!isEnabled ? 'opacity-50' : ''}
-              />
+                  onChange={(newValue: unknown) => {
+                    const option = newValue as { value: string; provider: string } | null;
+                    if (option) {
+                      if (option.value === '__custom__') {
+                        setIsLeadCustomModel(true);
+                        setLeadModel('');
+                        return;
+                      }
+                      setLeadModel(option.value);
+                      setLeadProvider(option.provider);
+                    }
+                  }}
+                  placeholder="Select lead model..."
+                  isDisabled={!isEnabled}
+                  className={!isEnabled ? 'opacity-50' : ''}
+                />
+              ) : (
+                <Input
+                  className="h-[38px] mb-2"
+                  placeholder="Type model name here"
+                  onChange={(event) => setLeadModel(event.target.value)}
+                  value={leadModel}
+                  disabled={!isEnabled}
+                />
+              )}
               <p className={`text-xs ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
                 Strong model for initial planning and fallback recovery
               </p>
             </div>
 
             <div className="space-y-2">
-              <label className={`text-sm ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
-                Worker Model
-              </label>
-              <Select
-                options={modelOptions}
-                value={modelOptions.find((opt) => opt.value === workerModel) || null}
-                onChange={(newValue: unknown) => {
-                  const option = newValue as { value: string; provider: string } | null;
-                  if (option) {
-                    setWorkerModel(option.value);
-                    setWorkerProvider(option.provider);
+              <div className="flex items-center justify-between">
+                <label className={`text-sm ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
+                  Worker Model
+                </label>
+                {isWorkerCustomModel && (
+                  <button
+                    onClick={() => setIsWorkerCustomModel(false)}
+                    className={`text-xs ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'} hover:underline`}
+                    type="button"
+                  >
+                    Back to model list
+                  </button>
+                )}
+              </div>
+              {!isWorkerCustomModel ? (
+                <Select
+                  options={modelOptions}
+                  value={
+                    workerModel
+                      ? modelOptions.find((opt) => opt.value === workerModel) || null
+                      : null
                   }
-                }}
-                placeholder="Select worker model..."
-                isDisabled={!isEnabled}
-                className={!isEnabled ? 'opacity-50' : ''}
-              />
+                  onChange={(newValue: unknown) => {
+                    const option = newValue as { value: string; provider: string } | null;
+                    if (option) {
+                      if (option.value === '__custom__') {
+                        setIsWorkerCustomModel(true);
+                        setWorkerModel('');
+                        return;
+                      }
+                      setWorkerModel(option.value);
+                      setWorkerProvider(option.provider);
+                    }
+                  }}
+                  placeholder="Select worker model..."
+                  isDisabled={!isEnabled}
+                  className={!isEnabled ? 'opacity-50' : ''}
+                />
+              ) : (
+                <Input
+                  className="h-[38px] mb-2"
+                  placeholder="Type model name here"
+                  onChange={(event) => setWorkerModel(event.target.value)}
+                  value={workerModel}
+                  disabled={!isEnabled}
+                />
+              )}
               <p className={`text-xs ${!isEnabled ? 'text-text-muted' : 'text-textSubtle'}`}>
                 Fast model for routine execution tasks
               </p>