custom.rs

  1use crate::{AgentServer, AgentServerDelegate, load_proxy_env};
  2use acp_thread::AgentConnection;
  3use agent_client_protocol as acp;
  4use anyhow::{Context as _, Result};
  5use collections::HashSet;
  6use fs::Fs;
  7use gpui::{App, AppContext as _, SharedString, Task};
  8use project::agent_server_store::{AllAgentServersSettings, ExternalAgentServerName};
  9use settings::{SettingsStore, update_settings_file};
 10use std::{path::Path, rc::Rc, sync::Arc};
 11use ui::IconName;
 12
 13/// A generic agent server implementation for custom user-defined agents
 14pub struct CustomAgentServer {
 15    name: SharedString,
 16}
 17
 18impl CustomAgentServer {
 19    pub fn new(name: SharedString) -> Self {
 20        Self { name }
 21    }
 22}
 23
 24impl AgentServer for CustomAgentServer {
 25    fn name(&self) -> SharedString {
 26        self.name.clone()
 27    }
 28
 29    fn logo(&self) -> IconName {
 30        IconName::Terminal
 31    }
 32
 33    fn default_mode(&self, cx: &App) -> Option<acp::SessionModeId> {
 34        let settings = cx.read_global(|settings: &SettingsStore, _| {
 35            settings
 36                .get::<AllAgentServersSettings>(None)
 37                .custom
 38                .get(&self.name())
 39                .cloned()
 40        });
 41
 42        settings
 43            .as_ref()
 44            .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
 45    }
 46
 47    fn favorite_config_option_value_ids(
 48        &self,
 49        config_id: &acp::SessionConfigId,
 50        cx: &mut App,
 51    ) -> HashSet<acp::SessionConfigValueId> {
 52        let settings = cx.read_global(|settings: &SettingsStore, _| {
 53            settings
 54                .get::<AllAgentServersSettings>(None)
 55                .custom
 56                .get(&self.name())
 57                .cloned()
 58        });
 59
 60        settings
 61            .as_ref()
 62            .and_then(|s| s.favorite_config_option_values(config_id.0.as_ref()))
 63            .map(|values| {
 64                values
 65                    .iter()
 66                    .cloned()
 67                    .map(acp::SessionConfigValueId::new)
 68                    .collect()
 69            })
 70            .unwrap_or_default()
 71    }
 72
 73    fn toggle_favorite_config_option_value(
 74        &self,
 75        config_id: acp::SessionConfigId,
 76        value_id: acp::SessionConfigValueId,
 77        should_be_favorite: bool,
 78        fs: Arc<dyn Fs>,
 79        cx: &App,
 80    ) {
 81        let name = self.name();
 82        let config_id = config_id.to_string();
 83        let value_id = value_id.to_string();
 84
 85        update_settings_file(fs, cx, move |settings, _| {
 86            let settings = settings
 87                .agent_servers
 88                .get_or_insert_default()
 89                .custom
 90                .entry(name.clone())
 91                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
 92                    default_model: None,
 93                    default_mode: None,
 94                    favorite_models: Vec::new(),
 95                    default_config_options: Default::default(),
 96                    favorite_config_option_values: Default::default(),
 97                });
 98
 99            match settings {
100                settings::CustomAgentServerSettings::Custom {
101                    favorite_config_option_values,
102                    ..
103                }
104                | settings::CustomAgentServerSettings::Extension {
105                    favorite_config_option_values,
106                    ..
107                } => {
108                    let entry = favorite_config_option_values
109                        .entry(config_id.clone())
110                        .or_insert_with(Vec::new);
111
112                    if should_be_favorite {
113                        if !entry.iter().any(|v| v == &value_id) {
114                            entry.push(value_id.clone());
115                        }
116                    } else {
117                        entry.retain(|v| v != &value_id);
118                        if entry.is_empty() {
119                            favorite_config_option_values.remove(&config_id);
120                        }
121                    }
122                }
123            }
124        });
125    }
126
127    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
128        let name = self.name();
129        update_settings_file(fs, cx, move |settings, _| {
130            let settings = settings
131                .agent_servers
132                .get_or_insert_default()
133                .custom
134                .entry(name.clone())
135                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
136                    default_model: None,
137                    default_mode: None,
138                    favorite_models: Vec::new(),
139                    default_config_options: Default::default(),
140                    favorite_config_option_values: Default::default(),
141                });
142
143            match settings {
144                settings::CustomAgentServerSettings::Custom { default_mode, .. }
145                | settings::CustomAgentServerSettings::Extension { default_mode, .. } => {
146                    *default_mode = mode_id.map(|m| m.to_string());
147                }
148            }
149        });
150    }
151
152    fn default_model(&self, cx: &App) -> Option<acp::ModelId> {
153        let settings = cx.read_global(|settings: &SettingsStore, _| {
154            settings
155                .get::<AllAgentServersSettings>(None)
156                .custom
157                .get(&self.name())
158                .cloned()
159        });
160
161        settings
162            .as_ref()
163            .and_then(|s| s.default_model().map(acp::ModelId::new))
164    }
165
166    fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
167        let name = self.name();
168        update_settings_file(fs, cx, move |settings, _| {
169            let settings = settings
170                .agent_servers
171                .get_or_insert_default()
172                .custom
173                .entry(name.clone())
174                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
175                    default_model: None,
176                    default_mode: None,
177                    favorite_models: Vec::new(),
178                    default_config_options: Default::default(),
179                    favorite_config_option_values: Default::default(),
180                });
181
182            match settings {
183                settings::CustomAgentServerSettings::Custom { default_model, .. }
184                | settings::CustomAgentServerSettings::Extension { default_model, .. } => {
185                    *default_model = model_id.map(|m| m.to_string());
186                }
187            }
188        });
189    }
190
191    fn favorite_model_ids(&self, cx: &mut App) -> HashSet<acp::ModelId> {
192        let settings = cx.read_global(|settings: &SettingsStore, _| {
193            settings
194                .get::<AllAgentServersSettings>(None)
195                .custom
196                .get(&self.name())
197                .cloned()
198        });
199
200        settings
201            .as_ref()
202            .map(|s| {
203                s.favorite_models()
204                    .iter()
205                    .map(|id| acp::ModelId::new(id.clone()))
206                    .collect()
207            })
208            .unwrap_or_default()
209    }
210
211    fn toggle_favorite_model(
212        &self,
213        model_id: acp::ModelId,
214        should_be_favorite: bool,
215        fs: Arc<dyn Fs>,
216        cx: &App,
217    ) {
218        let name = self.name();
219        update_settings_file(fs, cx, move |settings, _| {
220            let settings = settings
221                .agent_servers
222                .get_or_insert_default()
223                .custom
224                .entry(name.clone())
225                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
226                    default_model: None,
227                    default_mode: None,
228                    favorite_models: Vec::new(),
229                    default_config_options: Default::default(),
230                    favorite_config_option_values: Default::default(),
231                });
232
233            let favorite_models = match settings {
234                settings::CustomAgentServerSettings::Custom {
235                    favorite_models, ..
236                }
237                | settings::CustomAgentServerSettings::Extension {
238                    favorite_models, ..
239                } => favorite_models,
240            };
241
242            let model_id_str = model_id.to_string();
243            if should_be_favorite {
244                if !favorite_models.contains(&model_id_str) {
245                    favorite_models.push(model_id_str);
246                }
247            } else {
248                favorite_models.retain(|id| id != &model_id_str);
249            }
250        });
251    }
252
253    fn default_config_option(&self, config_id: &str, cx: &App) -> Option<String> {
254        let settings = cx.read_global(|settings: &SettingsStore, _| {
255            settings
256                .get::<AllAgentServersSettings>(None)
257                .custom
258                .get(&self.name())
259                .cloned()
260        });
261
262        settings
263            .as_ref()
264            .and_then(|s| s.default_config_option(config_id).map(|s| s.to_string()))
265    }
266
267    fn set_default_config_option(
268        &self,
269        config_id: &str,
270        value_id: Option<&str>,
271        fs: Arc<dyn Fs>,
272        cx: &mut App,
273    ) {
274        let name = self.name();
275        let config_id = config_id.to_string();
276        let value_id = value_id.map(|s| s.to_string());
277        update_settings_file(fs, cx, move |settings, _| {
278            let settings = settings
279                .agent_servers
280                .get_or_insert_default()
281                .custom
282                .entry(name.clone())
283                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
284                    default_model: None,
285                    default_mode: None,
286                    favorite_models: Vec::new(),
287                    default_config_options: Default::default(),
288                    favorite_config_option_values: Default::default(),
289                });
290
291            match settings {
292                settings::CustomAgentServerSettings::Custom {
293                    default_config_options,
294                    ..
295                }
296                | settings::CustomAgentServerSettings::Extension {
297                    default_config_options,
298                    ..
299                } => {
300                    if let Some(value) = value_id.clone() {
301                        default_config_options.insert(config_id.clone(), value);
302                    } else {
303                        default_config_options.remove(&config_id);
304                    }
305                }
306            }
307        });
308    }
309
310    fn connect(
311        &self,
312        root_dir: Option<&Path>,
313        delegate: AgentServerDelegate,
314        cx: &mut App,
315    ) -> Task<Result<(Rc<dyn AgentConnection>, Option<task::SpawnInTerminal>)>> {
316        let name = self.name();
317        let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
318        let is_remote = delegate.project.read(cx).is_via_remote_server();
319        let default_mode = self.default_mode(cx);
320        let default_model = self.default_model(cx);
321        let default_config_options = cx.read_global(|settings: &SettingsStore, _| {
322            settings
323                .get::<AllAgentServersSettings>(None)
324                .custom
325                .get(&self.name())
326                .map(|s| match s {
327                    project::agent_server_store::CustomAgentServerSettings::Custom {
328                        default_config_options,
329                        ..
330                    }
331                    | project::agent_server_store::CustomAgentServerSettings::Extension {
332                        default_config_options,
333                        ..
334                    } => default_config_options.clone(),
335                })
336                .unwrap_or_default()
337        });
338        let store = delegate.store.downgrade();
339        let extra_env = load_proxy_env(cx);
340        cx.spawn(async move |cx| {
341            let (command, root_dir, login) = store
342                .update(cx, |store, cx| {
343                    let agent = store
344                        .get_external_agent(&ExternalAgentServerName(name.clone()))
345                        .with_context(|| {
346                            format!("Custom agent server `{}` is not registered", name)
347                        })?;
348                    anyhow::Ok(agent.get_command(
349                        root_dir.as_deref(),
350                        extra_env,
351                        delegate.status_tx,
352                        delegate.new_version_available,
353                        &mut cx.to_async(),
354                    ))
355                })??
356                .await?;
357            let connection = crate::acp::connect(
358                name,
359                command,
360                root_dir.as_ref(),
361                default_mode,
362                default_model,
363                default_config_options,
364                is_remote,
365                cx,
366            )
367            .await?;
368            Ok((connection, login))
369        })
370    }
371
372    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
373        self
374    }
375}