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().as_ref())
 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().as_ref())
 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.to_string())
 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                | settings::CustomAgentServerSettings::Registry {
109                    favorite_config_option_values,
110                    ..
111                } => {
112                    let entry = favorite_config_option_values
113                        .entry(config_id.clone())
114                        .or_insert_with(Vec::new);
115
116                    if should_be_favorite {
117                        if !entry.iter().any(|v| v == &value_id) {
118                            entry.push(value_id.clone());
119                        }
120                    } else {
121                        entry.retain(|v| v != &value_id);
122                        if entry.is_empty() {
123                            favorite_config_option_values.remove(&config_id);
124                        }
125                    }
126                }
127            }
128        });
129    }
130
131    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
132        let name = self.name();
133        update_settings_file(fs, cx, move |settings, _| {
134            let settings = settings
135                .agent_servers
136                .get_or_insert_default()
137                .custom
138                .entry(name.to_string())
139                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
140                    default_model: None,
141                    default_mode: None,
142                    favorite_models: Vec::new(),
143                    default_config_options: Default::default(),
144                    favorite_config_option_values: Default::default(),
145                });
146
147            match settings {
148                settings::CustomAgentServerSettings::Custom { default_mode, .. }
149                | settings::CustomAgentServerSettings::Extension { default_mode, .. }
150                | settings::CustomAgentServerSettings::Registry { default_mode, .. } => {
151                    *default_mode = mode_id.map(|m| m.to_string());
152                }
153            }
154        });
155    }
156
157    fn default_model(&self, cx: &App) -> Option<acp::ModelId> {
158        let settings = cx.read_global(|settings: &SettingsStore, _| {
159            settings
160                .get::<AllAgentServersSettings>(None)
161                .custom
162                .get(self.name().as_ref())
163                .cloned()
164        });
165
166        settings
167            .as_ref()
168            .and_then(|s| s.default_model().map(acp::ModelId::new))
169    }
170
171    fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
172        let name = self.name();
173        update_settings_file(fs, cx, move |settings, _| {
174            let settings = settings
175                .agent_servers
176                .get_or_insert_default()
177                .custom
178                .entry(name.to_string())
179                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
180                    default_model: None,
181                    default_mode: None,
182                    favorite_models: Vec::new(),
183                    default_config_options: Default::default(),
184                    favorite_config_option_values: Default::default(),
185                });
186
187            match settings {
188                settings::CustomAgentServerSettings::Custom { default_model, .. }
189                | settings::CustomAgentServerSettings::Extension { default_model, .. }
190                | settings::CustomAgentServerSettings::Registry { default_model, .. } => {
191                    *default_model = model_id.map(|m| m.to_string());
192                }
193            }
194        });
195    }
196
197    fn favorite_model_ids(&self, cx: &mut App) -> HashSet<acp::ModelId> {
198        let settings = cx.read_global(|settings: &SettingsStore, _| {
199            settings
200                .get::<AllAgentServersSettings>(None)
201                .custom
202                .get(self.name().as_ref())
203                .cloned()
204        });
205
206        settings
207            .as_ref()
208            .map(|s| {
209                s.favorite_models()
210                    .iter()
211                    .map(|id| acp::ModelId::new(id.clone()))
212                    .collect()
213            })
214            .unwrap_or_default()
215    }
216
217    fn toggle_favorite_model(
218        &self,
219        model_id: acp::ModelId,
220        should_be_favorite: bool,
221        fs: Arc<dyn Fs>,
222        cx: &App,
223    ) {
224        let name = self.name();
225        update_settings_file(fs, cx, move |settings, _| {
226            let settings = settings
227                .agent_servers
228                .get_or_insert_default()
229                .custom
230                .entry(name.to_string())
231                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
232                    default_model: None,
233                    default_mode: None,
234                    favorite_models: Vec::new(),
235                    default_config_options: Default::default(),
236                    favorite_config_option_values: Default::default(),
237                });
238
239            let favorite_models = match settings {
240                settings::CustomAgentServerSettings::Custom {
241                    favorite_models, ..
242                }
243                | settings::CustomAgentServerSettings::Extension {
244                    favorite_models, ..
245                }
246                | settings::CustomAgentServerSettings::Registry {
247                    favorite_models, ..
248                } => favorite_models,
249            };
250
251            let model_id_str = model_id.to_string();
252            if should_be_favorite {
253                if !favorite_models.contains(&model_id_str) {
254                    favorite_models.push(model_id_str);
255                }
256            } else {
257                favorite_models.retain(|id| id != &model_id_str);
258            }
259        });
260    }
261
262    fn default_config_option(&self, config_id: &str, cx: &App) -> Option<String> {
263        let settings = cx.read_global(|settings: &SettingsStore, _| {
264            settings
265                .get::<AllAgentServersSettings>(None)
266                .custom
267                .get(self.name().as_ref())
268                .cloned()
269        });
270
271        settings
272            .as_ref()
273            .and_then(|s| s.default_config_option(config_id).map(|s| s.to_string()))
274    }
275
276    fn set_default_config_option(
277        &self,
278        config_id: &str,
279        value_id: Option<&str>,
280        fs: Arc<dyn Fs>,
281        cx: &mut App,
282    ) {
283        let name = self.name();
284        let config_id = config_id.to_string();
285        let value_id = value_id.map(|s| s.to_string());
286        update_settings_file(fs, cx, move |settings, _| {
287            let settings = settings
288                .agent_servers
289                .get_or_insert_default()
290                .custom
291                .entry(name.to_string())
292                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
293                    default_model: None,
294                    default_mode: None,
295                    favorite_models: Vec::new(),
296                    default_config_options: Default::default(),
297                    favorite_config_option_values: Default::default(),
298                });
299
300            match settings {
301                settings::CustomAgentServerSettings::Custom {
302                    default_config_options,
303                    ..
304                }
305                | settings::CustomAgentServerSettings::Extension {
306                    default_config_options,
307                    ..
308                }
309                | settings::CustomAgentServerSettings::Registry {
310                    default_config_options,
311                    ..
312                } => {
313                    if let Some(value) = value_id.clone() {
314                        default_config_options.insert(config_id.clone(), value);
315                    } else {
316                        default_config_options.remove(&config_id);
317                    }
318                }
319            }
320        });
321    }
322
323    fn connect(
324        &self,
325        root_dir: Option<&Path>,
326        delegate: AgentServerDelegate,
327        cx: &mut App,
328    ) -> Task<Result<(Rc<dyn AgentConnection>, Option<task::SpawnInTerminal>)>> {
329        let name = self.name();
330        let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
331        let is_remote = delegate.project.read(cx).is_via_remote_server();
332        let default_mode = self.default_mode(cx);
333        let default_model = self.default_model(cx);
334        let default_config_options = cx.read_global(|settings: &SettingsStore, _| {
335            settings
336                .get::<AllAgentServersSettings>(None)
337                .custom
338                .get(self.name().as_ref())
339                .map(|s| match s {
340                    project::agent_server_store::CustomAgentServerSettings::Custom {
341                        default_config_options,
342                        ..
343                    }
344                    | project::agent_server_store::CustomAgentServerSettings::Extension {
345                        default_config_options,
346                        ..
347                    }
348                    | project::agent_server_store::CustomAgentServerSettings::Registry {
349                        default_config_options,
350                        ..
351                    } => default_config_options.clone(),
352                })
353                .unwrap_or_default()
354        });
355        let store = delegate.store.downgrade();
356        let extra_env = load_proxy_env(cx);
357        cx.spawn(async move |cx| {
358            let (command, root_dir, login) = store
359                .update(cx, |store, cx| {
360                    let agent = store
361                        .get_external_agent(&ExternalAgentServerName(name.clone()))
362                        .with_context(|| {
363                            format!("Custom agent server `{}` is not registered", name)
364                        })?;
365                    anyhow::Ok(agent.get_command(
366                        root_dir.as_deref(),
367                        extra_env,
368                        delegate.status_tx,
369                        delegate.new_version_available,
370                        &mut cx.to_async(),
371                    ))
372                })??
373                .await?;
374            let connection = crate::acp::connect(
375                name,
376                command,
377                root_dir.as_ref(),
378                default_mode,
379                default_model,
380                default_config_options,
381                is_remote,
382                cx,
383            )
384            .await?;
385            Ok((connection, login))
386        })
387    }
388
389    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
390        self
391    }
392}