model_selector.rs

  1use std::{cmp::Reverse, rc::Rc, sync::Arc};
  2
  3use acp_thread::{AgentModelIcon, AgentModelInfo, AgentModelList, AgentModelSelector};
  4use agent_client_protocol::ModelId;
  5use agent_servers::AgentServer;
  6use agent_settings::AgentSettings;
  7use anyhow::Result;
  8use collections::{HashSet, IndexMap};
  9use fs::Fs;
 10use futures::FutureExt;
 11use fuzzy::{StringMatchCandidate, match_strings};
 12use gpui::{
 13    Action, AsyncWindowContext, BackgroundExecutor, DismissEvent, FocusHandle, Subscription, Task,
 14    WeakEntity,
 15};
 16use itertools::Itertools;
 17use ordered_float::OrderedFloat;
 18use picker::{Picker, PickerDelegate};
 19use settings::{Settings, SettingsStore};
 20use ui::{DocumentationAside, DocumentationSide, IntoElement, prelude::*};
 21use util::ResultExt;
 22use zed_actions::agent::OpenSettings;
 23
 24use crate::ui::{HoldForDefault, ModelSelectorFooter, ModelSelectorHeader, ModelSelectorListItem};
 25
 26pub type AcpModelSelector = Picker<AcpModelPickerDelegate>;
 27
 28pub fn acp_model_selector(
 29    selector: Rc<dyn AgentModelSelector>,
 30    agent_server: Rc<dyn AgentServer>,
 31    fs: Arc<dyn Fs>,
 32    focus_handle: FocusHandle,
 33    window: &mut Window,
 34    cx: &mut Context<AcpModelSelector>,
 35) -> AcpModelSelector {
 36    let delegate =
 37        AcpModelPickerDelegate::new(selector, agent_server, fs, focus_handle, window, cx);
 38    Picker::list(delegate, window, cx)
 39        .show_scrollbar(true)
 40        .width(rems(20.))
 41        .max_height(Some(rems(20.).into()))
 42}
 43
 44enum AcpModelPickerEntry {
 45    Separator(SharedString),
 46    Model(AgentModelInfo, bool),
 47}
 48
 49pub struct AcpModelPickerDelegate {
 50    selector: Rc<dyn AgentModelSelector>,
 51    agent_server: Rc<dyn AgentServer>,
 52    fs: Arc<dyn Fs>,
 53    filtered_entries: Vec<AcpModelPickerEntry>,
 54    models: Option<AgentModelList>,
 55    selected_index: usize,
 56    selected_description: Option<(usize, SharedString, bool)>,
 57    selected_model: Option<AgentModelInfo>,
 58    favorites: HashSet<ModelId>,
 59    _refresh_models_task: Task<()>,
 60    _settings_subscription: Subscription,
 61    focus_handle: FocusHandle,
 62}
 63
 64impl AcpModelPickerDelegate {
 65    fn new(
 66        selector: Rc<dyn AgentModelSelector>,
 67        agent_server: Rc<dyn AgentServer>,
 68        fs: Arc<dyn Fs>,
 69        focus_handle: FocusHandle,
 70        window: &mut Window,
 71        cx: &mut Context<AcpModelSelector>,
 72    ) -> Self {
 73        let rx = selector.watch(cx);
 74        let refresh_models_task = {
 75            cx.spawn_in(window, {
 76                async move |this, cx| {
 77                    async fn refresh(
 78                        this: &WeakEntity<Picker<AcpModelPickerDelegate>>,
 79                        cx: &mut AsyncWindowContext,
 80                    ) -> Result<()> {
 81                        let (models_task, selected_model_task) = this.update(cx, |this, cx| {
 82                            (
 83                                this.delegate.selector.list_models(cx),
 84                                this.delegate.selector.selected_model(cx),
 85                            )
 86                        })?;
 87
 88                        let (models, selected_model) =
 89                            futures::join!(models_task, selected_model_task);
 90
 91                        this.update_in(cx, |this, window, cx| {
 92                            this.delegate.models = models.ok();
 93                            this.delegate.selected_model = selected_model.ok();
 94                            this.refresh(window, cx)
 95                        })
 96                    }
 97
 98                    refresh(&this, cx).await.log_err();
 99                    if let Some(mut rx) = rx {
100                        while let Ok(()) = rx.recv().await {
101                            refresh(&this, cx).await.log_err();
102                        }
103                    }
104                }
105            })
106        };
107
108        let agent_server_for_subscription = agent_server.clone();
109        let settings_subscription =
110            cx.observe_global_in::<SettingsStore>(window, move |picker, window, cx| {
111                // Only refresh if the favorites actually changed to avoid redundant work
112                // when other settings are modified (e.g., user editing settings.json)
113                let new_favorites = agent_server_for_subscription.favorite_model_ids(cx);
114                if new_favorites != picker.delegate.favorites {
115                    picker.delegate.favorites = new_favorites;
116                    picker.refresh(window, cx);
117                }
118            });
119        let favorites = agent_server.favorite_model_ids(cx);
120
121        Self {
122            selector,
123            agent_server,
124            fs,
125            filtered_entries: Vec::new(),
126            models: None,
127            selected_model: None,
128            selected_index: 0,
129            selected_description: None,
130            favorites,
131            _refresh_models_task: refresh_models_task,
132            _settings_subscription: settings_subscription,
133            focus_handle,
134        }
135    }
136
137    pub fn active_model(&self) -> Option<&AgentModelInfo> {
138        self.selected_model.as_ref()
139    }
140
141    pub fn favorites_count(&self) -> usize {
142        self.favorites.len()
143    }
144
145    pub fn cycle_favorite_models(&mut self, window: &mut Window, cx: &mut Context<Picker<Self>>) {
146        if self.favorites.is_empty() {
147            return;
148        }
149
150        let Some(models) = &self.models else {
151            return;
152        };
153
154        let all_models: Vec<&AgentModelInfo> = match models {
155            AgentModelList::Flat(list) => list.iter().collect(),
156            AgentModelList::Grouped(index_map) => index_map.values().flatten().collect(),
157        };
158
159        let favorite_models: Vec<_> = all_models
160            .into_iter()
161            .filter(|model| self.favorites.contains(&model.id))
162            .unique_by(|model| &model.id)
163            .collect();
164
165        if favorite_models.is_empty() {
166            return;
167        }
168
169        let current_id = self.selected_model.as_ref().map(|m| &m.id);
170
171        let current_index_in_favorites = current_id
172            .and_then(|id| favorite_models.iter().position(|m| &m.id == id))
173            .unwrap_or(usize::MAX);
174
175        let next_index = if current_index_in_favorites == usize::MAX {
176            0
177        } else {
178            (current_index_in_favorites + 1) % favorite_models.len()
179        };
180
181        let next_model = favorite_models[next_index].clone();
182
183        self.selector
184            .select_model(next_model.id.clone(), cx)
185            .detach_and_log_err(cx);
186
187        self.selected_model = Some(next_model);
188
189        // Keep the picker selection aligned with the newly-selected model
190        if let Some(new_index) = self.filtered_entries.iter().position(|entry| {
191            matches!(entry, AcpModelPickerEntry::Model(model_info, _) if self.selected_model.as_ref().is_some_and(|selected| model_info.id == selected.id))
192        }) {
193            self.set_selected_index(new_index, window, cx);
194        } else {
195            cx.notify();
196        }
197    }
198}
199
200impl PickerDelegate for AcpModelPickerDelegate {
201    type ListItem = AnyElement;
202
203    fn match_count(&self) -> usize {
204        self.filtered_entries.len()
205    }
206
207    fn selected_index(&self) -> usize {
208        self.selected_index
209    }
210
211    fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context<Picker<Self>>) {
212        self.selected_index = ix.min(self.filtered_entries.len().saturating_sub(1));
213        cx.notify();
214    }
215
216    fn can_select(
217        &mut self,
218        ix: usize,
219        _window: &mut Window,
220        _cx: &mut Context<Picker<Self>>,
221    ) -> bool {
222        match self.filtered_entries.get(ix) {
223            Some(AcpModelPickerEntry::Model(_, _)) => true,
224            Some(AcpModelPickerEntry::Separator(_)) | None => false,
225        }
226    }
227
228    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
229        "Select a model…".into()
230    }
231
232    fn update_matches(
233        &mut self,
234        query: String,
235        window: &mut Window,
236        cx: &mut Context<Picker<Self>>,
237    ) -> Task<()> {
238        let favorites = self.favorites.clone();
239
240        cx.spawn_in(window, async move |this, cx| {
241            let filtered_models = match this
242                .read_with(cx, |this, cx| {
243                    this.delegate.models.clone().map(move |models| {
244                        fuzzy_search(models, query, cx.background_executor().clone())
245                    })
246                })
247                .ok()
248                .flatten()
249            {
250                Some(task) => task.await,
251                None => AgentModelList::Flat(vec![]),
252            };
253
254            this.update_in(cx, |this, window, cx| {
255                this.delegate.filtered_entries =
256                    info_list_to_picker_entries(filtered_models, &favorites);
257                // Finds the currently selected model in the list
258                let new_index = this
259                    .delegate
260                    .selected_model
261                    .as_ref()
262                    .and_then(|selected| {
263                        this.delegate.filtered_entries.iter().position(|entry| {
264                            if let AcpModelPickerEntry::Model(model_info, _) = entry {
265                                model_info.id == selected.id
266                            } else {
267                                false
268                            }
269                        })
270                    })
271                    .unwrap_or(0);
272                this.set_selected_index(new_index, Some(picker::Direction::Down), true, window, cx);
273                cx.notify();
274            })
275            .ok();
276        })
277    }
278
279    fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
280        if let Some(AcpModelPickerEntry::Model(model_info, _)) =
281            self.filtered_entries.get(self.selected_index)
282        {
283            if window.modifiers().secondary() {
284                let default_model = self.agent_server.default_model(cx);
285                let is_default = default_model.as_ref() == Some(&model_info.id);
286
287                self.agent_server.set_default_model(
288                    if is_default {
289                        None
290                    } else {
291                        Some(model_info.id.clone())
292                    },
293                    self.fs.clone(),
294                    cx,
295                );
296            }
297
298            self.selector
299                .select_model(model_info.id.clone(), cx)
300                .detach_and_log_err(cx);
301            self.selected_model = Some(model_info.clone());
302            let current_index = self.selected_index;
303            self.set_selected_index(current_index, window, cx);
304
305            cx.emit(DismissEvent);
306        }
307    }
308
309    fn dismissed(&mut self, window: &mut Window, cx: &mut Context<Picker<Self>>) {
310        cx.defer_in(window, |picker, window, cx| {
311            picker.set_query("", window, cx);
312        });
313    }
314
315    fn render_match(
316        &self,
317        ix: usize,
318        selected: bool,
319        _: &mut Window,
320        cx: &mut Context<Picker<Self>>,
321    ) -> Option<Self::ListItem> {
322        match self.filtered_entries.get(ix)? {
323            AcpModelPickerEntry::Separator(title) => {
324                Some(ModelSelectorHeader::new(title, ix > 1).into_any_element())
325            }
326            AcpModelPickerEntry::Model(model_info, is_favorite) => {
327                let is_selected = Some(model_info) == self.selected_model.as_ref();
328                let default_model = self.agent_server.default_model(cx);
329                let is_default = default_model.as_ref() == Some(&model_info.id);
330
331                let is_favorite = *is_favorite;
332                let handle_action_click = {
333                    let model_id = model_info.id.clone();
334                    let fs = self.fs.clone();
335                    let agent_server = self.agent_server.clone();
336
337                    cx.listener(move |_, _, _, cx| {
338                        agent_server.toggle_favorite_model(
339                            model_id.clone(),
340                            !is_favorite,
341                            fs.clone(),
342                            cx,
343                        );
344                    })
345                };
346
347                Some(
348                    div()
349                        .id(("model-picker-menu-child", ix))
350                        .when_some(model_info.description.clone(), |this, description| {
351                            this.on_hover(cx.listener(move |menu, hovered, _, cx| {
352                                if *hovered {
353                                    menu.delegate.selected_description =
354                                        Some((ix, description.clone(), is_default));
355                                } else if matches!(menu.delegate.selected_description, Some((id, _, _)) if id == ix) {
356                                    menu.delegate.selected_description = None;
357                                }
358                                cx.notify();
359                            }))
360                        })
361                        .child(
362                            ModelSelectorListItem::new(ix, model_info.name.clone())
363                                .map(|this| match &model_info.icon {
364                                    Some(AgentModelIcon::Path(path)) => this.icon_path(path.clone()),
365                                    Some(AgentModelIcon::Named(icon)) => this.icon(*icon),
366                                    None => this,
367                                })
368                                .is_selected(is_selected)
369                                .is_focused(selected)
370                                .is_favorite(is_favorite)
371                                .on_toggle_favorite(handle_action_click),
372                        )
373                        .into_any_element(),
374                )
375            }
376        }
377    }
378
379    fn documentation_aside(
380        &self,
381        _window: &mut Window,
382        cx: &mut Context<Picker<Self>>,
383    ) -> Option<ui::DocumentationAside> {
384        self.selected_description
385            .as_ref()
386            .map(|(_, description, is_default)| {
387                let description = description.clone();
388                let is_default = *is_default;
389
390                let settings = AgentSettings::get_global(cx);
391                let side = match settings.dock {
392                    settings::DockPosition::Left => DocumentationSide::Right,
393                    settings::DockPosition::Bottom | settings::DockPosition::Right => {
394                        DocumentationSide::Left
395                    }
396                };
397
398                DocumentationAside::new(
399                    side,
400                    Rc::new(move |_| {
401                        v_flex()
402                            .gap_1()
403                            .child(Label::new(description.clone()))
404                            .child(HoldForDefault::new(is_default))
405                            .into_any_element()
406                    }),
407                )
408            })
409    }
410
411    fn documentation_aside_index(&self) -> Option<usize> {
412        self.selected_description.as_ref().map(|(ix, _, _)| *ix)
413    }
414
415    fn render_footer(
416        &self,
417        _window: &mut Window,
418        _cx: &mut Context<Picker<Self>>,
419    ) -> Option<AnyElement> {
420        let focus_handle = self.focus_handle.clone();
421
422        if !self.selector.should_render_footer() {
423            return None;
424        }
425
426        Some(ModelSelectorFooter::new(OpenSettings.boxed_clone(), focus_handle).into_any_element())
427    }
428}
429
430fn info_list_to_picker_entries(
431    model_list: AgentModelList,
432    favorites: &HashSet<ModelId>,
433) -> Vec<AcpModelPickerEntry> {
434    let mut entries = Vec::new();
435
436    let all_models: Vec<_> = match &model_list {
437        AgentModelList::Flat(list) => list.iter().collect(),
438        AgentModelList::Grouped(index_map) => index_map.values().flatten().collect(),
439    };
440
441    let favorite_models: Vec<_> = all_models
442        .iter()
443        .filter(|m| favorites.contains(&m.id))
444        .unique_by(|m| &m.id)
445        .collect();
446
447    let has_favorites = !favorite_models.is_empty();
448    if has_favorites {
449        entries.push(AcpModelPickerEntry::Separator("Favorite".into()));
450        for model in favorite_models {
451            entries.push(AcpModelPickerEntry::Model((*model).clone(), true));
452        }
453    }
454
455    match model_list {
456        AgentModelList::Flat(list) => {
457            if has_favorites {
458                entries.push(AcpModelPickerEntry::Separator("All".into()));
459            }
460            for model in list {
461                let is_favorite = favorites.contains(&model.id);
462                entries.push(AcpModelPickerEntry::Model(model, is_favorite));
463            }
464        }
465        AgentModelList::Grouped(index_map) => {
466            for (group_name, models) in index_map {
467                entries.push(AcpModelPickerEntry::Separator(group_name.0));
468                for model in models {
469                    let is_favorite = favorites.contains(&model.id);
470                    entries.push(AcpModelPickerEntry::Model(model, is_favorite));
471                }
472            }
473        }
474    }
475
476    entries
477}
478
479async fn fuzzy_search(
480    model_list: AgentModelList,
481    query: String,
482    executor: BackgroundExecutor,
483) -> AgentModelList {
484    async fn fuzzy_search_list(
485        model_list: Vec<AgentModelInfo>,
486        query: &str,
487        executor: BackgroundExecutor,
488    ) -> Vec<AgentModelInfo> {
489        let candidates = model_list
490            .iter()
491            .enumerate()
492            .map(|(ix, model)| StringMatchCandidate::new(ix, model.name.as_ref()))
493            .collect::<Vec<_>>();
494        let mut matches = match_strings(
495            &candidates,
496            query,
497            false,
498            true,
499            100,
500            &Default::default(),
501            executor,
502        )
503        .await;
504
505        matches.sort_unstable_by_key(|mat| {
506            let candidate = &candidates[mat.candidate_id];
507            (Reverse(OrderedFloat(mat.score)), candidate.id)
508        });
509
510        matches
511            .into_iter()
512            .map(|mat| model_list[mat.candidate_id].clone())
513            .collect()
514    }
515
516    match model_list {
517        AgentModelList::Flat(model_list) => {
518            AgentModelList::Flat(fuzzy_search_list(model_list, &query, executor).await)
519        }
520        AgentModelList::Grouped(index_map) => {
521            let groups =
522                futures::future::join_all(index_map.into_iter().map(|(group_name, models)| {
523                    fuzzy_search_list(models, &query, executor.clone())
524                        .map(|results| (group_name, results))
525                }))
526                .await;
527            AgentModelList::Grouped(IndexMap::from_iter(
528                groups
529                    .into_iter()
530                    .filter(|(_, results)| !results.is_empty()),
531            ))
532        }
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use agent_client_protocol as acp;
539    use gpui::TestAppContext;
540
541    use super::*;
542
543    fn create_model_list(grouped_models: Vec<(&str, Vec<&str>)>) -> AgentModelList {
544        AgentModelList::Grouped(IndexMap::from_iter(grouped_models.into_iter().map(
545            |(group, models)| {
546                (
547                    acp_thread::AgentModelGroupName(group.to_string().into()),
548                    models
549                        .into_iter()
550                        .map(|model| acp_thread::AgentModelInfo {
551                            id: acp::ModelId::new(model.to_string()),
552                            name: model.to_string().into(),
553                            description: None,
554                            icon: None,
555                        })
556                        .collect::<Vec<_>>(),
557                )
558            },
559        )))
560    }
561
562    fn assert_models_eq(result: AgentModelList, expected: Vec<(&str, Vec<&str>)>) {
563        let AgentModelList::Grouped(groups) = result else {
564            panic!("Expected LanguageModelInfoList::Grouped, got {:?}", result);
565        };
566
567        assert_eq!(
568            groups.len(),
569            expected.len(),
570            "Number of groups doesn't match"
571        );
572
573        for (i, (expected_group, expected_models)) in expected.iter().enumerate() {
574            let (actual_group, actual_models) = groups.get_index(i).unwrap();
575            assert_eq!(
576                actual_group.0.as_ref(),
577                *expected_group,
578                "Group at position {} doesn't match expected group",
579                i
580            );
581            assert_eq!(
582                actual_models.len(),
583                expected_models.len(),
584                "Number of models in group {} doesn't match",
585                expected_group
586            );
587
588            for (j, expected_model_name) in expected_models.iter().enumerate() {
589                assert_eq!(
590                    actual_models[j].name, *expected_model_name,
591                    "Model at position {} in group {} doesn't match expected model",
592                    j, expected_group
593                );
594            }
595        }
596    }
597
598    fn create_favorites(models: Vec<&str>) -> HashSet<ModelId> {
599        models
600            .into_iter()
601            .map(|m| ModelId::new(m.to_string()))
602            .collect()
603    }
604
605    fn get_entry_model_ids(entries: &[AcpModelPickerEntry]) -> Vec<&str> {
606        entries
607            .iter()
608            .filter_map(|entry| match entry {
609                AcpModelPickerEntry::Model(info, _) => Some(info.id.0.as_ref()),
610                _ => None,
611            })
612            .collect()
613    }
614
615    fn get_entry_labels(entries: &[AcpModelPickerEntry]) -> Vec<&str> {
616        entries
617            .iter()
618            .map(|entry| match entry {
619                AcpModelPickerEntry::Model(info, _) => info.id.0.as_ref(),
620                AcpModelPickerEntry::Separator(s) => &s,
621            })
622            .collect()
623    }
624
625    #[gpui::test]
626    async fn test_fuzzy_match(cx: &mut TestAppContext) {
627        let models = create_model_list(vec![
628            (
629                "zed",
630                vec![
631                    "Claude 3.7 Sonnet",
632                    "Claude 3.7 Sonnet Thinking",
633                    "gpt-4.1",
634                    "gpt-4.1-nano",
635                ],
636            ),
637            ("openai", vec!["gpt-3.5-turbo", "gpt-4.1", "gpt-4.1-nano"]),
638            ("ollama", vec!["mistral", "deepseek"]),
639        ]);
640
641        // Results should preserve models order whenever possible.
642        // In the case below, `zed/gpt-4.1` and `openai/gpt-4.1` have identical
643        // similarity scores, but `zed/gpt-4.1` was higher in the models list,
644        // so it should appear first in the results.
645        let results = fuzzy_search(models.clone(), "41".into(), cx.executor()).await;
646        assert_models_eq(
647            results,
648            vec![
649                ("zed", vec!["gpt-4.1", "gpt-4.1-nano"]),
650                ("openai", vec!["gpt-4.1", "gpt-4.1-nano"]),
651            ],
652        );
653
654        // Fuzzy search
655        let results = fuzzy_search(models.clone(), "4n".into(), cx.executor()).await;
656        assert_models_eq(
657            results,
658            vec![
659                ("zed", vec!["gpt-4.1-nano"]),
660                ("openai", vec!["gpt-4.1-nano"]),
661            ],
662        );
663    }
664
665    #[gpui::test]
666    fn test_favorites_section_appears_when_favorites_exist(_cx: &mut TestAppContext) {
667        let models = create_model_list(vec![
668            ("zed", vec!["zed/claude", "zed/gemini"]),
669            ("openai", vec!["openai/gpt-5"]),
670        ]);
671        let favorites = create_favorites(vec!["zed/gemini"]);
672
673        let entries = info_list_to_picker_entries(models, &favorites);
674
675        assert!(matches!(
676            entries.first(),
677            Some(AcpModelPickerEntry::Separator(s)) if s == "Favorite"
678        ));
679
680        let model_ids = get_entry_model_ids(&entries);
681        assert_eq!(model_ids[0], "zed/gemini");
682    }
683
684    #[gpui::test]
685    fn test_no_favorites_section_when_no_favorites(_cx: &mut TestAppContext) {
686        let models = create_model_list(vec![("zed", vec!["zed/claude", "zed/gemini"])]);
687        let favorites = create_favorites(vec![]);
688
689        let entries = info_list_to_picker_entries(models, &favorites);
690
691        assert!(matches!(
692            entries.first(),
693            Some(AcpModelPickerEntry::Separator(s)) if s == "zed"
694        ));
695    }
696
697    #[gpui::test]
698    fn test_models_have_correct_actions(_cx: &mut TestAppContext) {
699        let models = create_model_list(vec![
700            ("zed", vec!["zed/claude", "zed/gemini"]),
701            ("openai", vec!["openai/gpt-5"]),
702        ]);
703        let favorites = create_favorites(vec!["zed/claude"]);
704
705        let entries = info_list_to_picker_entries(models, &favorites);
706
707        for entry in &entries {
708            if let AcpModelPickerEntry::Model(info, is_favorite) = entry {
709                if info.id.0.as_ref() == "zed/claude" {
710                    assert!(is_favorite, "zed/claude should be a favorite");
711                } else {
712                    assert!(!is_favorite, "{} should not be a favorite", info.id.0);
713                }
714            }
715        }
716    }
717
718    #[gpui::test]
719    fn test_favorites_appear_in_both_sections(_cx: &mut TestAppContext) {
720        let models = create_model_list(vec![
721            ("zed", vec!["zed/claude", "zed/gemini"]),
722            ("openai", vec!["openai/gpt-5", "openai/gpt-4"]),
723        ]);
724        let favorites = create_favorites(vec!["zed/gemini", "openai/gpt-5"]);
725
726        let entries = info_list_to_picker_entries(models, &favorites);
727        let model_ids = get_entry_model_ids(&entries);
728
729        assert_eq!(model_ids[0], "zed/gemini");
730        assert_eq!(model_ids[1], "openai/gpt-5");
731
732        assert!(model_ids[2..].contains(&"zed/gemini"));
733        assert!(model_ids[2..].contains(&"openai/gpt-5"));
734    }
735
736    #[gpui::test]
737    fn test_favorites_are_not_duplicated_when_repeated_in_other_sections(_cx: &mut TestAppContext) {
738        let models = create_model_list(vec![
739            ("Recommended", vec!["zed/claude", "anthropic/claude"]),
740            ("Zed", vec!["zed/claude", "zed/gpt-5"]),
741            ("Antropic", vec!["anthropic/claude"]),
742            ("OpenAI", vec!["openai/gpt-5"]),
743        ]);
744
745        let favorites = create_favorites(vec!["zed/claude"]);
746
747        let entries = info_list_to_picker_entries(models, &favorites);
748        let labels = get_entry_labels(&entries);
749
750        assert_eq!(
751            labels,
752            vec![
753                "Favorite",
754                "zed/claude",
755                "Recommended",
756                "zed/claude",
757                "anthropic/claude",
758                "Zed",
759                "zed/claude",
760                "zed/gpt-5",
761                "Antropic",
762                "anthropic/claude",
763                "OpenAI",
764                "openai/gpt-5"
765            ]
766        );
767    }
768
769    #[gpui::test]
770    fn test_flat_model_list_with_favorites(_cx: &mut TestAppContext) {
771        let models = AgentModelList::Flat(vec![
772            acp_thread::AgentModelInfo {
773                id: acp::ModelId::new("zed/claude".to_string()),
774                name: "Claude".into(),
775                description: None,
776                icon: None,
777            },
778            acp_thread::AgentModelInfo {
779                id: acp::ModelId::new("zed/gemini".to_string()),
780                name: "Gemini".into(),
781                description: None,
782                icon: None,
783            },
784        ]);
785        let favorites = create_favorites(vec!["zed/gemini"]);
786
787        let entries = info_list_to_picker_entries(models, &favorites);
788
789        assert!(matches!(
790            entries.first(),
791            Some(AcpModelPickerEntry::Separator(s)) if s == "Favorite"
792        ));
793
794        assert!(entries.iter().any(|e| matches!(
795            e,
796            AcpModelPickerEntry::Separator(s) if s == "All"
797        )));
798    }
799
800    #[gpui::test]
801    fn test_favorites_count_returns_correct_count(_cx: &mut TestAppContext) {
802        let empty_favorites: HashSet<ModelId> = HashSet::default();
803        assert_eq!(empty_favorites.len(), 0);
804
805        let one_favorite = create_favorites(vec!["model-a"]);
806        assert_eq!(one_favorite.len(), 1);
807
808        let multiple_favorites = create_favorites(vec!["model-a", "model-b", "model-c"]);
809        assert_eq!(multiple_favorites.len(), 3);
810
811        let with_duplicates = create_favorites(vec!["model-a", "model-a", "model-b"]);
812        assert_eq!(with_duplicates.len(), 2);
813    }
814
815    #[gpui::test]
816    fn test_is_favorite_flag_set_correctly_in_entries(_cx: &mut TestAppContext) {
817        let models = AgentModelList::Flat(vec![
818            acp_thread::AgentModelInfo {
819                id: acp::ModelId::new("favorite-model".to_string()),
820                name: "Favorite".into(),
821                description: None,
822                icon: None,
823            },
824            acp_thread::AgentModelInfo {
825                id: acp::ModelId::new("regular-model".to_string()),
826                name: "Regular".into(),
827                description: None,
828                icon: None,
829            },
830        ]);
831        let favorites = create_favorites(vec!["favorite-model"]);
832
833        let entries = info_list_to_picker_entries(models, &favorites);
834
835        for entry in &entries {
836            if let AcpModelPickerEntry::Model(info, is_favorite) = entry {
837                if info.id.0.as_ref() == "favorite-model" {
838                    assert!(*is_favorite, "favorite-model should have is_favorite=true");
839                } else if info.id.0.as_ref() == "regular-model" {
840                    assert!(!*is_favorite, "regular-model should have is_favorite=false");
841                }
842            }
843        }
844    }
845}