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