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: &mut App) -> Option<acp::SessionModeId> {
34 let settings = cx.read_global(|settings: &SettingsStore, _| {
35 settings
36 .get::<AllAgentServersSettings>(None)
37 .custom
38 .get(&self.name())
39 .cloned()
40 });
41
42 settings
43 .as_ref()
44 .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
45 }
46
47 fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
48 let name = self.name();
49 update_settings_file(fs, cx, move |settings, _| {
50 let settings = settings
51 .agent_servers
52 .get_or_insert_default()
53 .custom
54 .entry(name.clone())
55 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
56 default_model: None,
57 default_mode: None,
58 favorite_models: Vec::new(),
59 });
60
61 match settings {
62 settings::CustomAgentServerSettings::Custom { default_mode, .. }
63 | settings::CustomAgentServerSettings::Extension { default_mode, .. } => {
64 *default_mode = mode_id.map(|m| m.to_string());
65 }
66 }
67 });
68 }
69
70 fn default_model(&self, cx: &mut App) -> Option<acp::ModelId> {
71 let settings = cx.read_global(|settings: &SettingsStore, _| {
72 settings
73 .get::<AllAgentServersSettings>(None)
74 .custom
75 .get(&self.name())
76 .cloned()
77 });
78
79 settings
80 .as_ref()
81 .and_then(|s| s.default_model().map(acp::ModelId::new))
82 }
83
84 fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
85 let name = self.name();
86 update_settings_file(fs, cx, move |settings, _| {
87 let settings = settings
88 .agent_servers
89 .get_or_insert_default()
90 .custom
91 .entry(name.clone())
92 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
93 default_model: None,
94 default_mode: None,
95 favorite_models: Vec::new(),
96 });
97
98 match settings {
99 settings::CustomAgentServerSettings::Custom { default_model, .. }
100 | settings::CustomAgentServerSettings::Extension { default_model, .. } => {
101 *default_model = model_id.map(|m| m.to_string());
102 }
103 }
104 });
105 }
106
107 fn favorite_model_ids(&self, cx: &mut App) -> HashSet<acp::ModelId> {
108 let settings = cx.read_global(|settings: &SettingsStore, _| {
109 settings
110 .get::<AllAgentServersSettings>(None)
111 .custom
112 .get(&self.name())
113 .cloned()
114 });
115
116 settings
117 .as_ref()
118 .map(|s| {
119 s.favorite_models()
120 .iter()
121 .map(|id| acp::ModelId::new(id.clone()))
122 .collect()
123 })
124 .unwrap_or_default()
125 }
126
127 fn toggle_favorite_model(
128 &self,
129 model_id: acp::ModelId,
130 should_be_favorite: bool,
131 fs: Arc<dyn Fs>,
132 cx: &App,
133 ) {
134 let name = self.name();
135 update_settings_file(fs, cx, move |settings, _| {
136 let settings = settings
137 .agent_servers
138 .get_or_insert_default()
139 .custom
140 .entry(name.clone())
141 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
142 default_model: None,
143 default_mode: None,
144 favorite_models: Vec::new(),
145 });
146
147 let favorite_models = match settings {
148 settings::CustomAgentServerSettings::Custom {
149 favorite_models, ..
150 }
151 | settings::CustomAgentServerSettings::Extension {
152 favorite_models, ..
153 } => favorite_models,
154 };
155
156 let model_id_str = model_id.to_string();
157 if should_be_favorite {
158 if !favorite_models.contains(&model_id_str) {
159 favorite_models.push(model_id_str);
160 }
161 } else {
162 favorite_models.retain(|id| id != &model_id_str);
163 }
164 });
165 }
166
167 fn connect(
168 &self,
169 root_dir: Option<&Path>,
170 delegate: AgentServerDelegate,
171 cx: &mut App,
172 ) -> Task<Result<(Rc<dyn AgentConnection>, Option<task::SpawnInTerminal>)>> {
173 let name = self.name();
174 let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
175 let is_remote = delegate.project.read(cx).is_via_remote_server();
176 let default_mode = self.default_mode(cx);
177 let default_model = self.default_model(cx);
178 let store = delegate.store.downgrade();
179 let extra_env = load_proxy_env(cx);
180 cx.spawn(async move |cx| {
181 let (command, root_dir, login) = store
182 .update(cx, |store, cx| {
183 let agent = store
184 .get_external_agent(&ExternalAgentServerName(name.clone()))
185 .with_context(|| {
186 format!("Custom agent server `{}` is not registered", name)
187 })?;
188 anyhow::Ok(agent.get_command(
189 root_dir.as_deref(),
190 extra_env,
191 delegate.status_tx,
192 delegate.new_version_available,
193 &mut cx.to_async(),
194 ))
195 })??
196 .await?;
197 let connection = crate::acp::connect(
198 name,
199 command,
200 root_dir.as_ref(),
201 default_mode,
202 default_model,
203 is_remote,
204 cx,
205 )
206 .await?;
207 Ok((connection, login))
208 })
209 }
210
211 fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
212 self
213 }
214}