1use crate::{AgentServerDelegate, load_proxy_env};
  2use acp_thread::AgentConnection;
  3use agent_client_protocol as acp;
  4use anyhow::{Context as _, Result};
  5use collections::HashMap;
  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 crate::AgentServer for CustomAgentServer {
 25    fn telemetry_id(&self) -> &'static str {
 26        "custom"
 27    }
 28
 29    fn name(&self) -> SharedString {
 30        self.name.clone()
 31    }
 32
 33    fn logo(&self) -> IconName {
 34        IconName::Terminal
 35    }
 36
 37    fn default_mode(&self, cx: &mut App) -> Option<acp::SessionModeId> {
 38        let settings = cx.read_global(|settings: &SettingsStore, _| {
 39            settings
 40                .get::<AllAgentServersSettings>(None)
 41                .custom
 42                .get(&self.name())
 43                .cloned()
 44        });
 45
 46        settings
 47            .as_ref()
 48            .and_then(|s| s.default_mode.clone().map(|m| acp::SessionModeId(m.into())))
 49    }
 50
 51    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
 52        let name = self.name();
 53        update_settings_file(fs, cx, move |settings, _| {
 54            settings
 55                .agent_servers
 56                .get_or_insert_default()
 57                .custom
 58                .get_mut(&name)
 59                .unwrap()
 60                .default_mode = mode_id.map(|m| m.to_string())
 61        });
 62    }
 63
 64    fn connect(
 65        &self,
 66        root_dir: Option<&Path>,
 67        delegate: AgentServerDelegate,
 68        cx: &mut App,
 69    ) -> Task<
 70        Result<(
 71            Rc<dyn AgentConnection>,
 72            HashMap<String, task::SpawnInTerminal>,
 73        )>,
 74    > {
 75        let name = self.name();
 76        let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
 77        let is_remote = delegate.project.read(cx).is_via_remote_server();
 78        let default_mode = self.default_mode(cx);
 79        let store = delegate.store.downgrade();
 80        let extra_env = load_proxy_env(cx);
 81
 82        cx.spawn(async move |cx| {
 83            let (command, root_dir, auth_commands) = store
 84                .update(cx, |store, cx| {
 85                    let agent = store
 86                        .get_external_agent(&ExternalAgentServerName(name.clone()))
 87                        .with_context(|| {
 88                            format!("Custom agent server `{}` is not registered", name)
 89                        })?;
 90                    anyhow::Ok(agent.get_command(
 91                        root_dir.as_deref(),
 92                        extra_env,
 93                        delegate.status_tx,
 94                        delegate.new_version_available,
 95                        &mut cx.to_async(),
 96                    ))
 97                })??
 98                .await?;
 99            let connection = crate::acp::connect(
100                name,
101                command,
102                root_dir.as_ref(),
103                default_mode,
104                is_remote,
105                cx,
106            )
107            .await?;
108            Ok((connection, auth_commands))
109        })
110    }
111
112    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
113        self
114    }
115}