From 899dfd314934b49fd152942a3d2cd8fb07263323 Mon Sep 17 00:00:00 2001
From: Zane Staggs <zane@squareup.com>
Date: Sun, 13 Jul 2025 00:17:37 -0700
Subject: [PATCH] fix lead/worker display and selection

---
 .../models/bottom_bar/ModelsBottomBar.tsx     | 53 +++++++++++++++++--
 .../subcomponents/LeadWorkerSettings.tsx      | 24 +++++++--
 ui/desktop/src/components/ui/Select.tsx       | 22 ++++----
 3 files changed, 79 insertions(+), 20 deletions(-)

diff --git a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
index 610fc2eca0..68ddd737f0 100644
--- a/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
+++ b/ui/desktop/src/components/settings/models/bottom_bar/ModelsBottomBar.tsx
@@ -63,16 +63,62 @@ export default function ModelsBottomBar({
         const leadModel = await read('GOOSE_LEAD_MODEL', false);
         setIsLeadWorkerActive(!!leadModel);
       } catch (error) {
+        console.error('Error checking lead model:', error);
         setIsLeadWorkerActive(false);
       }
     };
     checkLeadWorker();
   }, [read]);
 
+  // Refresh lead/worker status when modal closes
+  const handleLeadWorkerModalClose = () => {
+    setIsLeadWorkerModalOpen(false);
+    // Refresh the lead/worker status after modal closes
+    const checkLeadWorker = async () => {
+      try {
+        const leadModel = await read('GOOSE_LEAD_MODEL', false);
+        const currentModel = await read('GOOSE_MODEL', false);
+        setIsLeadWorkerActive(!!leadModel);
+        setLeadModelName((leadModel as string) || '');
+        setCurrentActiveModel((currentModel as string) || '');
+      } catch (error) {
+        console.error('Error checking lead model after modal close:', error);
+        setIsLeadWorkerActive(false);
+      }
+    };
+    checkLeadWorker();
+  };
+
   // Determine which model to display - activeModel takes priority when lead/worker is active
   const displayModel =
     isLeadWorkerActive && currentModelInfo?.model ? currentModelInfo.model : displayModelName;
-  const modelMode = currentModelInfo?.mode;
+
+  // Since currentModelInfo.mode is not working, let's determine mode differently
+  // We'll need to get the lead model and compare it with the current model
+  const [leadModelName, setLeadModelName] = useState<string>('');
+  const [currentActiveModel, setCurrentActiveModel] = useState<string>('');
+
+  // Get lead model name and current model for comparison
+  useEffect(() => {
+    const getModelInfo = async () => {
+      try {
+        const leadModel = await read('GOOSE_LEAD_MODEL', false);
+        const currentModel = await read('GOOSE_MODEL', false);
+        setLeadModelName((leadModel as string) || '');
+        setCurrentActiveModel((currentModel as string) || '');
+      } catch (error) {
+        console.error('Error getting model info:', error);
+      }
+    };
+    getModelInfo();
+  }, [read]);
+
+  // Determine the mode based on which model is currently active
+  const modelMode = isLeadWorkerActive
+    ? currentActiveModel === leadModelName
+      ? 'lead'
+      : 'worker'
+    : undefined;
 
   // Update display provider when current provider changes
   useEffect(() => {
@@ -209,10 +255,7 @@ export default function ModelsBottomBar({
       ) : null}
 
       {isLeadWorkerModalOpen ? (
-        <LeadWorkerSettings
-          isOpen={isLeadWorkerModalOpen}
-          onClose={() => setIsLeadWorkerModalOpen(false)}
-        />
+        <LeadWorkerSettings isOpen={isLeadWorkerModalOpen} onClose={handleLeadWorkerModalClose} />
       ) : null}
 
       {/* Save Recipe Dialog - copied from RecipeEditor.tsx */}
diff --git a/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx b/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx
index 247d2db6f3..cecfc430b7 100644
--- a/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx
+++ b/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx
@@ -5,7 +5,7 @@ import { Button } from '../../../ui/button';
 import { Select } from '../../../ui/Select';
 import { Input } from '../../../ui/input';
 import { Info } from 'lucide-react';
-import { Dialog, DialogContent } from '../../../ui/dialog';
+import { Dialog, DialogContent, DialogHeader, DialogTitle } from '../../../ui/dialog';
 
 interface LeadWorkerSettingsProps {
   isOpen: boolean;
@@ -30,6 +30,8 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
 
   // Load current configuration
   useEffect(() => {
+    if (!isOpen) return; // Only load when modal is open
+
     const loadConfig = async () => {
       try {
         setIsLoading(true);
@@ -50,11 +52,18 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
         if (leadModelConfig) {
           setLeadModel(leadModelConfig as string);
           setIsEnabled(true);
+        } else {
+          setLeadModel('');
+          setIsEnabled(false);
         }
         if (leadProviderConfig) setLeadProvider(leadProviderConfig as string);
+        else setLeadProvider('');
         if (leadTurnsConfig) setLeadTurns(Number(leadTurnsConfig));
+        else setLeadTurns(3);
         if (failureThresholdConfig) setFailureThreshold(Number(failureThresholdConfig));
+        else setFailureThreshold(2);
         if (fallbackTurnsConfig) setFallbackTurns(Number(fallbackTurnsConfig));
+        else setFallbackTurns(2);
 
         // Set worker model to current model or from config
         const workerModelConfig = await read('GOOSE_MODEL', false);
@@ -62,11 +71,15 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
           setWorkerModel(workerModelConfig as string);
         } else if (currentModel) {
           setWorkerModel(currentModel as string);
+        } else {
+          setWorkerModel('');
         }
 
         const workerProviderConfig = await read('GOOSE_PROVIDER', false);
         if (workerProviderConfig) {
           setWorkerProvider(workerProviderConfig as string);
+        } else {
+          setWorkerProvider('');
         }
 
         // Load available models
@@ -95,7 +108,7 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
     };
 
     loadConfig();
-  }, [read, getProviders, currentModel]);
+  }, [read, getProviders, currentModel, isOpen]);
 
   const handleSave = async () => {
     try {
@@ -130,6 +143,9 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
     return (
       <Dialog open={isOpen} onOpenChange={(open) => !open && onClose()}>
         <DialogContent className="sm:max-w-[500px]">
+          <DialogHeader>
+            <DialogTitle>Lead/Worker Mode</DialogTitle>
+          </DialogHeader>
           <div className="p-4">Loading...</div>
         </DialogContent>
       </Dialog>
@@ -139,9 +155,11 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
   return (
     <Dialog open={isOpen} onOpenChange={(open) => !open && onClose()}>
       <DialogContent className="sm:max-w-[500px]">
+        <DialogHeader>
+          <DialogTitle>Lead/Worker Mode</DialogTitle>
+        </DialogHeader>
         <div className="p-4 space-y-4">
           <div className="space-y-2">
-            <h3 className="text-lg font-medium text-textProminent">Lead/Worker Mode</h3>
             <p className="text-sm text-textSubtle">
               Configure a lead model for planning and a worker model for execution
             </p>
diff --git a/ui/desktop/src/components/ui/Select.tsx b/ui/desktop/src/components/ui/Select.tsx
index a23e5e20c6..cd868a74a5 100644
--- a/ui/desktop/src/components/ui/Select.tsx
+++ b/ui/desktop/src/components/ui/Select.tsx
@@ -9,17 +9,14 @@ export const Select = (props: React.ComponentProps<typeof ReactSelect>) => {
       isSearchable={props.isSearchable !== false}
       closeMenuOnSelect={props.closeMenuOnSelect !== false}
       blurInputOnSelect={props.blurInputOnSelect !== false}
-      menuPortalTarget={document.body}
-      menuPosition="fixed"
-      menuPlacement="auto"
       classNames={{
         container: () => 'w-full cursor-pointer relative',
         indicatorSeparator: () => 'h-0',
         control: ({ isFocused }) =>
           `border ${isFocused ? 'border-borderStandard' : 'border-borderSubtle'} focus:border-borderStandard hover:border-borderStandard rounded-md w-full px-4 py-2 text-sm text-textSubtle hover:cursor-pointer`,
         menu: () =>
-          'mt-1 bg-background-default border border-borderStandard rounded-md text-textSubtle overflow-hidden shadow-lg select__menu',
-        menuPortal: () => 'z-[9999]',
+          'mt-1 bg-background-default border border-borderStandard rounded-md text-textSubtle shadow-lg select__menu z-[9999] absolute',
+        menuList: () => 'max-h-60 overflow-y-auto py-1',
         option: ({ isFocused, isSelected, isDisabled }) => {
           let classes = 'py-2 px-4 text-sm cursor-pointer';
 
@@ -37,19 +34,20 @@ export const Select = (props: React.ComponentProps<typeof ReactSelect>) => {
         },
       }}
       menuShouldBlockScroll={false}
-      menuShouldScrollIntoView={true}
+      menuShouldScrollIntoView={false}
       tabSelectsValue={true}
       openMenuOnFocus={false}
-      captureMenuScroll={false}
       styles={{
-        menuPortal: (base) => ({
-          ...base,
-          zIndex: 9999,
-          pointerEvents: 'auto',
-        }),
         menu: (base) => ({
           ...base,
           pointerEvents: 'auto',
+          zIndex: 9999,
+        }),
+        menuList: (base) => ({
+          ...base,
+          maxHeight: '240px',
+          overflowY: 'auto',
+          pointerEvents: 'auto',
         }),
         option: (base) => ({
           ...base,