fix: switch model when selecting an agent (#1362)

* fix: switch model when selecting an agent

The early return in setAgent() checked if the current model was valid
and returned immediately without ever falling through to the agent's
configured model. This meant selecting an agent never updated the
model dropdown to match the agent's preferred model.

Removing the early return allows the fallback chain to work correctly:
1. Persisted session-specific agent-model selection
2. Settings default model
3. Agent's own model (previously unreachable)
4. No change (current model stays)

* fix: prefer agent model over saved selection

* fix: preserve restored model on agent hydration

---------
Co-authored-by: Bohdan Triapitsyn <artmore@protonmail.com>

packages/ui/src/components/chat/ModelControls.tsx

@@ -618,6 +618,7 @@ export const ModelControls: React.FC<ModelControlsProps> = ({
     ];
 
     const prevAgentNameRef = React.useRef<string | undefined>(undefined);
+    const explicitAgentSwitchRef = React.useRef<string | null>(null);
     const latestLoadedUserChoiceRestoreRef = React.useRef<string | null>(null);
 
     const currentSessionDirectory = currentSessionId ? getDirectoryForSession(currentSessionId) : undefined;
@@ -1007,6 +1008,9 @@ export const ModelControls: React.FC<ModelControlsProps> = ({
                     prevAgentNameRef.current = currentAgentName;
 
                     if (currentAgentName && currentSessionId) {
+                        const shouldPreferAgentModel = explicitAgentSwitchRef.current === currentAgentName;
+                        explicitAgentSwitchRef.current = null;
+
                         await new Promise<void>((resolve) => {
                             const timer = setTimeout(resolve, 50);
                             abortController.signal.addEventListener('abort', () => {
@@ -1019,6 +1023,33 @@ export const ModelControls: React.FC<ModelControlsProps> = ({
                             return;
                         }
 
+                        const selectedAgent = shouldPreferAgentModel
+                            ? agents.find((agent) => agent.name === currentAgentName)
+                            : undefined;
+                        if (selectedAgent?.model?.providerID && selectedAgent.model.modelID) {
+                            const result = tryApplyModelSelection(
+                                selectedAgent.model.providerID,
+                                selectedAgent.model.modelID,
+                                currentAgentName,
+                            );
+                            if (result === 'applied' || result === 'provider-missing') {
+                                if (result === 'applied') {
+                                    saveSessionModelSelection(
+                                        currentSessionId,
+                                        selectedAgent.model.providerID,
+                                        selectedAgent.model.modelID,
+                                    );
+                                    saveAgentModelForSession(
+                                        currentSessionId,
+                                        currentAgentName,
+                                        selectedAgent.model.providerID,
+                                        selectedAgent.model.modelID,
+                                    );
+                                }
+                                return;
+                            }
+                        }
+
                         const persistedChoice = getAgentModelForSession(currentSessionId, currentAgentName);
 
                         if (persistedChoice) {
@@ -1043,7 +1074,16 @@ export const ModelControls: React.FC<ModelControlsProps> = ({
         return () => {
             abortController.abort();
         };
-    }, [currentAgentName, currentSessionId, getAgentModelForSession, tryApplyModelSelection, contextHydrated]);
+    }, [
+        agents,
+        currentAgentName,
+        currentSessionId,
+        getAgentModelForSession,
+        saveAgentModelForSession,
+        saveSessionModelSelection,
+        tryApplyModelSelection,
+        contextHydrated,
+    ]);
 
     React.useEffect(() => {
         if (!contextHydrated || !currentAgentName) {
@@ -1119,6 +1159,7 @@ export const ModelControls: React.FC<ModelControlsProps> = ({
 
     const handleAgentChange = React.useCallback((agentName: string, options?: { closeModelSelector?: boolean }) => {
         try {
+            explicitAgentSwitchRef.current = agentName;
             setAgent(agentName);
             addRecentAgent(agentName);
             if (options?.closeModelSelector ?? true) {

packages/ui/src/stores/useConfigStore.ts

@@ -1739,6 +1739,20 @@ export const useConfigStore = create<ConfigStore>()(
                             });
                         };
 
+                        // Prefer the selected agent's configured model when switching agents.
+                        const agent = agents.find((candidate) => candidate.name === agentName);
+                        const agentModelSelection = agent?.model;
+                        if (agentModelSelection?.providerID && agentModelSelection?.modelID) {
+                            const { providerID, modelID } = agentModelSelection;
+                            const agentProvider = providers.find((provider) => provider.id === providerID);
+                            const agentModel = agentProvider?.models.find((model) => model.id === modelID);
+
+                            if (agentModel) {
+                                applyResolvedModelSelection(providerID, modelID, undefined);
+                                return;
+                            }
+                        }
+
                         if (currentSessionId) {
                             const existingAgentModel = useSelectionStore.getState().getAgentModelForSession(currentSessionId, agentName);
                             if (existingAgentModel && hasProviderModel(providers, existingAgentModel.providerId, existingAgentModel.modelId)) {
@@ -1759,11 +1773,7 @@ export const useConfigStore = create<ConfigStore>()(
                             }
                         }
 
-                        if (hasProviderModel(providers, currentProviderId, currentModelId)) {
+                        // If the agent has no preferred model, use settings default.
-                            return;
-                        }
-
-                        // If settings has a default model, use it instead of agent's preferred
                         if (settingsDefaultModel) {
                             const parsed = parseModelString(settingsDefaultModel);
                             if (parsed) {
@@ -1784,18 +1794,7 @@ export const useConfigStore = create<ConfigStore>()(
                             }
                         }
 
-                        // Fall back to agent's preferred model
+                        // Otherwise keep the current valid model selection unchanged.
-                        const agent = agents.find((candidate) => candidate.name === agentName);
-                        const agentModelSelection = agent?.model;
-                        if (agentModelSelection?.providerID && agentModelSelection?.modelID) {
-                            const { providerID, modelID } = agentModelSelection;
-                            const agentProvider = providers.find((provider) => provider.id === providerID);
-                            const agentModel = agentProvider?.models.find((model) => model.id === modelID);
-
-                            if (agentModel) {
-                                applyResolvedModelSelection(providerID, modelID, undefined);
-                            }
-                        }
                     }
                 },