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