1use crate::{AgentServer, AgentServerDelegate, load_proxy_env};
  2use acp_thread::AgentConnection;
  3use agent_client_protocol as acp;
  4use anyhow::{Context as _, Result};
  5use fs::Fs;
  6use gpui::{App, AppContext as _, SharedString, Task};
  7use project::agent_server_store::{AllAgentServersSettings, ExternalAgentServerName};
  8use settings::{SettingsStore, update_settings_file};
  9use std::{path::Path, rc::Rc, sync::Arc};
 10use ui::IconName;
 11
 12/// A generic agent server implementation for custom user-defined agents
 13pub struct CustomAgentServer {
 14    name: SharedString,
 15}
 16
 17impl CustomAgentServer {
 18    pub fn new(name: SharedString) -> Self {
 19        Self { name }
 20    }
 21}
 22
 23impl AgentServer for CustomAgentServer {
 24    fn name(&self) -> SharedString {
 25        self.name.clone()
 26    }
 27
 28    fn logo(&self) -> IconName {
 29        IconName::Terminal
 30    }
 31
 32    fn default_mode(&self, cx: &mut App) -> Option<acp::SessionModeId> {
 33        let settings = cx.read_global(|settings: &SettingsStore, _| {
 34            settings
 35                .get::<AllAgentServersSettings>(None)
 36                .custom
 37                .get(&self.name())
 38                .cloned()
 39        });
 40
 41        settings
 42            .as_ref()
 43            .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
 44    }
 45
 46    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
 47        let name = self.name();
 48        update_settings_file(fs, cx, move |settings, _| {
 49            let settings = settings
 50                .agent_servers
 51                .get_or_insert_default()
 52                .custom
 53                .entry(name.clone())
 54                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
 55                    default_model: None,
 56                    default_mode: None,
 57                });
 58
 59            match settings {
 60                settings::CustomAgentServerSettings::Custom { default_mode, .. }
 61                | settings::CustomAgentServerSettings::Extension { default_mode, .. } => {
 62                    *default_mode = mode_id.map(|m| m.to_string());
 63                }
 64            }
 65        });
 66    }
 67
 68    fn default_model(&self, cx: &mut App) -> Option<acp::ModelId> {
 69        let settings = cx.read_global(|settings: &SettingsStore, _| {
 70            settings
 71                .get::<AllAgentServersSettings>(None)
 72                .custom
 73                .get(&self.name())
 74                .cloned()
 75        });
 76
 77        settings
 78            .as_ref()
 79            .and_then(|s| s.default_model().map(acp::ModelId::new))
 80    }
 81
 82    fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
 83        let name = self.name();
 84        update_settings_file(fs, cx, move |settings, _| {
 85            let settings = settings
 86                .agent_servers
 87                .get_or_insert_default()
 88                .custom
 89                .entry(name.clone())
 90                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
 91                    default_model: None,
 92                    default_mode: None,
 93                });
 94
 95            match settings {
 96                settings::CustomAgentServerSettings::Custom { default_model, .. }
 97                | settings::CustomAgentServerSettings::Extension { default_model, .. } => {
 98                    *default_model = model_id.map(|m| m.to_string());
 99                }
100            }
101        });
102    }
103
104    fn connect(
105        &self,
106        root_dir: Option<&Path>,
107        delegate: AgentServerDelegate,
108        cx: &mut App,
109    ) -> Task<Result<(Rc<dyn AgentConnection>, Option<task::SpawnInTerminal>)>> {
110        let name = self.name();
111        let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
112        let is_remote = delegate.project.read(cx).is_via_remote_server();
113        let default_mode = self.default_mode(cx);
114        let default_model = self.default_model(cx);
115        let store = delegate.store.downgrade();
116        let extra_env = load_proxy_env(cx);
117        cx.spawn(async move |cx| {
118            let (command, root_dir, login) = store
119                .update(cx, |store, cx| {
120                    let agent = store
121                        .get_external_agent(&ExternalAgentServerName(name.clone()))
122                        .with_context(|| {
123                            format!("Custom agent server `{}` is not registered", name)
124                        })?;
125                    anyhow::Ok(agent.get_command(
126                        root_dir.as_deref(),
127                        extra_env,
128                        delegate.status_tx,
129                        delegate.new_version_available,
130                        &mut cx.to_async(),
131                    ))
132                })??
133                .await?;
134            let connection = crate::acp::connect(
135                name,
136                command,
137                root_dir.as_ref(),
138                default_mode,
139                default_model,
140                is_remote,
141                cx,
142            )
143            .await?;
144            Ok((connection, login))
145        })
146    }
147
148    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
149        self
150    }
151}