tenant.openagents/omega

1use acp_thread::AgentSessionModes;
2use agent_client_protocol::schema::v1 as acp;
3use agent_servers::AgentServer;
4
5use fs::Fs;
6use gpui::{Context, Entity, WeakEntity, Window, prelude::*};
7
8use std::{rc::Rc, sync::Arc};
9use ui::{
10    Button, ContextMenu, ContextMenuEntry, KeyBinding, PopoverMenu, PopoverMenuHandle, Tooltip,
11    prelude::*,
12};
13
14use crate::{CycleModeSelector, ToggleProfileSelector, ui::documentation_aside_side};
15
16pub struct ModeSelector {
17    connection: Rc<dyn AgentSessionModes>,
18    agent_server: Rc<dyn AgentServer>,
19    menu_handle: PopoverMenuHandle<ContextMenu>,
20    fs: Arc<dyn Fs>,
21    setting_mode: bool,
22}
23
24impl ModeSelector {
25    pub fn new(
26        session_modes: Rc<dyn AgentSessionModes>,
27        agent_server: Rc<dyn AgentServer>,
28        fs: Arc<dyn Fs>,
29    ) -> Self {
30        Self {
31            connection: session_modes,
32            agent_server,
33            menu_handle: PopoverMenuHandle::default(),
34            fs,
35            setting_mode: false,
36        }
37    }
38
39    pub fn menu_handle(&self) -> PopoverMenuHandle<ContextMenu> {
40        self.menu_handle.clone()
41    }
42
43    pub fn cycle_mode(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
44        let all_modes = self.connection.all_modes();
45        if all_modes.is_empty() {
46            return;
47        }
48
49        let current_mode = self.connection.current_mode();
50
51        let current_index = all_modes
52            .iter()
53            .position(|mode| mode.id.0 == current_mode.0)
54            .unwrap_or(0);
55
56        if let Some(next_mode) = all_modes.get((current_index + 1) % all_modes.len()) {
57            self.set_mode(next_mode.id.clone(), cx);
58        }
59    }
60
61    pub fn mode(&self) -> acp::SessionModeId {
62        self.connection.current_mode()
63    }
64
65    pub fn set_mode(&mut self, mode: acp::SessionModeId, cx: &mut Context<Self>) {
66        self.agent_server
67            .set_default_mode(Some(mode.clone()), self.fs.clone(), cx);
68
69        let task = self.connection.set_mode(mode, cx);
70        self.setting_mode = true;
71        cx.notify();
72
73        cx.spawn(async move |this: WeakEntity<ModeSelector>, cx| {
74            if let Err(err) = task.await {
75                log::error!("Failed to set session mode: {:?}", err);
76            }
77            this.update(cx, |this, cx| {
78                this.setting_mode = false;
79                cx.notify();
80            })
81            .ok();
82        })
83        .detach();
84    }
85
86    fn build_context_menu(
87        &self,
88        window: &mut Window,
89        cx: &mut Context<Self>,
90    ) -> Entity<ContextMenu> {
91        let weak_self = cx.weak_entity();
92
93        ContextMenu::build(window, cx, move |mut menu, _window, cx| {
94            let all_modes = self.connection.all_modes();
95            let current_mode = self.connection.current_mode();
96
97            let side = documentation_aside_side(cx);
98
99            for mode in all_modes {
100                let is_selected = &mode.id == &current_mode;
101                let entry = ContextMenuEntry::new(mode.name.clone())
102                    .toggleable(IconPosition::End, is_selected);
103
104                let entry = if let Some(description) = &mode.description {
105                    entry.documentation_aside(side, {
106                        let description = description.clone();
107
108                        move |_| Label::new(description.clone()).into_any_element()
109                    })
110                } else {
111                    entry
112                };
113
114                menu.push_item(entry.handler({
115                    let mode_id = mode.id.clone();
116                    let weak_self = weak_self.clone();
117                    move |_window, cx| {
118                        weak_self
119                            .update(cx, |this, cx| {
120                                this.set_mode(mode_id.clone(), cx);
121                            })
122                            .ok();
123                    }
124                }));
125            }
126
127            menu.key_context("ModeSelector")
128        })
129    }
130}
131
132impl Render for ModeSelector {
133    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
134        let current_mode_id = self.connection.current_mode();
135        let current_mode_name = self
136            .connection
137            .all_modes()
138            .iter()
139            .find(|mode| mode.id == current_mode_id)
140            .map(|mode| mode.name.clone())
141            .unwrap_or_else(|| "Unknown".into());
142
143        let this = cx.weak_entity();
144
145        let icon = if self.menu_handle.is_deployed() {
146            IconName::ChevronUp
147        } else {
148            IconName::ChevronDown
149        };
150
151        let trigger_button = Button::new("mode-selector-trigger", current_mode_name)
152            .label_size(LabelSize::Small)
153            .color(Color::Muted)
154            .end_icon(Icon::new(icon).size(IconSize::XSmall).color(Color::Muted))
155            .disabled(self.setting_mode);
156
157        PopoverMenu::new("mode-selector")
158            .trigger_with_tooltip(
159                trigger_button,
160                Tooltip::element({
161                    move |_window, cx| {
162                        v_flex()
163                            .gap_1()
164                            .child(
165                                h_flex()
166                                    .gap_2()
167                                    .justify_between()
168                                    .child(Label::new("Change Mode"))
169                                    .child(KeyBinding::for_action(&ToggleProfileSelector, cx)),
170                            )
171                            .child(
172                                h_flex()
173                                    .pt_1()
174                                    .gap_2()
175                                    .border_t_1()
176                                    .border_color(cx.theme().colors().border_variant)
177                                    .justify_between()
178                                    .child(Label::new("Cycle Through Modes"))
179                                    .child(KeyBinding::for_action(&CycleModeSelector, cx)),
180                            )
181                            .into_any()
182                    }
183                }),
184            )
185            .anchor(gpui::Anchor::BottomRight)
186            .with_handle(self.menu_handle.clone())
187            .offset(gpui::Point {
188                x: px(0.0),
189                y: px(-2.0),
190            })
191            .menu(move |window, cx| {
192                this.update(cx, |this, cx| this.build_context_menu(window, cx))
193                    .ok()
194            })
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use acp_thread::AgentConnection;
202    use fs::FakeFs;
203    use gpui::{App, Task, TestAppContext};
204    use parking_lot::Mutex;
205    use project::{AgentId, Project};
206    use std::{any::Any, cell::RefCell};
207
208    #[gpui::test]
209    fn setting_mode_saves_selected_mode_as_default(cx: &mut TestAppContext) {
210        let agent_server = Rc::new(TestAgentServer::default());
211        let session_modes = Rc::new(TestSessionModes::new());
212        let fs: Arc<dyn Fs> = FakeFs::new(cx.executor());
213
214        cx.update(|cx| {
215            let session_modes: Rc<dyn AgentSessionModes> = session_modes.clone();
216            let agent_server: Rc<dyn AgentServer> = agent_server.clone();
217            let selector = cx.new(|_| ModeSelector::new(session_modes, agent_server, fs));
218
219            selector.update(cx, |selector, cx| {
220                selector.set_mode(acp::SessionModeId::new("manual"), cx);
221            });
222        });
223
224        assert_eq!(
225            agent_server.saved_defaults.lock().as_slice(),
226            &[Some(acp::SessionModeId::new("manual"))]
227        );
228        assert_eq!(
229            session_modes.set_modes.borrow().as_slice(),
230            &[acp::SessionModeId::new("manual")]
231        );
232    }
233
234    #[derive(Default)]
235    struct TestAgentServer {
236        saved_defaults: Arc<Mutex<Vec<Option<acp::SessionModeId>>>>,
237    }
238
239    impl AgentServer for TestAgentServer {
240        fn logo(&self) -> IconName {
241            IconName::OmegaAssistant
242        }
243
244        fn agent_id(&self) -> AgentId {
245            AgentId::new("test-agent")
246        }
247
248        fn connect(
249            &self,
250            _delegate: agent_servers::AgentServerDelegate,
251            _project: Entity<Project>,
252            _cx: &mut App,
253        ) -> Task<anyhow::Result<Rc<dyn AgentConnection>>> {
254            Task::ready(Err(anyhow::anyhow!("test agent server cannot connect")))
255        }
256
257        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
258            self
259        }
260
261        fn set_default_mode(
262            &self,
263            mode_id: Option<acp::SessionModeId>,
264            _fs: Arc<dyn Fs>,
265            _cx: &mut App,
266        ) {
267            self.saved_defaults.lock().push(mode_id);
268        }
269    }
270
271    struct TestSessionModes {
272        current_mode: RefCell<acp::SessionModeId>,
273        set_modes: RefCell<Vec<acp::SessionModeId>>,
274    }
275
276    impl TestSessionModes {
277        fn new() -> Self {
278            Self {
279                current_mode: RefCell::new(acp::SessionModeId::new("auto")),
280                set_modes: RefCell::new(Vec::new()),
281            }
282        }
283    }
284
285    impl AgentSessionModes for TestSessionModes {
286        fn current_mode(&self) -> acp::SessionModeId {
287            self.current_mode.borrow().clone()
288        }
289
290        fn all_modes(&self) -> Vec<acp::SessionMode> {
291            vec![
292                acp::SessionMode::new("auto", "Auto"),
293                acp::SessionMode::new("manual", "Manual"),
294            ]
295        }
296
297        fn set_mode(&self, mode: acp::SessionModeId, _cx: &mut App) -> Task<anyhow::Result<()>> {
298            *self.current_mode.borrow_mut() = mode.clone();
299            self.set_modes.borrow_mut().push(mode);
300            Task::ready(Ok(()))
301        }
302    }
303}
304