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 credentials_provider::CredentialsProvider;
7use fs::Fs;
8use gpui::{App, AppContext as _, SharedString, Task};
9use language_model::{ApiKey, EnvVar};
10use project::agent_server_store::{
11 AllAgentServersSettings, CLAUDE_AGENT_NAME, CODEX_NAME, ExternalAgentServerName, GEMINI_NAME,
12};
13use settings::{SettingsStore, update_settings_file};
14use std::{rc::Rc, sync::Arc};
15use ui::IconName;
16
17/// A generic agent server implementation for custom user-defined agents
18pub struct CustomAgentServer {
19 name: SharedString,
20}
21
22impl CustomAgentServer {
23 pub fn new(name: SharedString) -> Self {
24 Self { name }
25 }
26}
27
28impl AgentServer for CustomAgentServer {
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: &App) -> Option<acp::SessionModeId> {
38 let settings = cx.read_global(|settings: &SettingsStore, _| {
39 settings
40 .get::<AllAgentServersSettings>(None)
41 .get(self.name().as_ref())
42 .cloned()
43 });
44
45 settings
46 .as_ref()
47 .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
48 }
49
50 fn favorite_config_option_value_ids(
51 &self,
52 config_id: &acp::SessionConfigId,
53 cx: &mut App,
54 ) -> HashSet<acp::SessionConfigValueId> {
55 let settings = cx.read_global(|settings: &SettingsStore, _| {
56 settings
57 .get::<AllAgentServersSettings>(None)
58 .get(self.name().as_ref())
59 .cloned()
60 });
61
62 settings
63 .as_ref()
64 .and_then(|s| s.favorite_config_option_values(config_id.0.as_ref()))
65 .map(|values| {
66 values
67 .iter()
68 .cloned()
69 .map(acp::SessionConfigValueId::new)
70 .collect()
71 })
72 .unwrap_or_default()
73 }
74
75 fn toggle_favorite_config_option_value(
76 &self,
77 config_id: acp::SessionConfigId,
78 value_id: acp::SessionConfigValueId,
79 should_be_favorite: bool,
80 fs: Arc<dyn Fs>,
81 cx: &App,
82 ) {
83 let name = self.name();
84 let config_id = config_id.to_string();
85 let value_id = value_id.to_string();
86
87 update_settings_file(fs, cx, move |settings, _| {
88 let settings = settings
89 .agent_servers
90 .get_or_insert_default()
91 .entry(name.to_string())
92 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
93 default_model: None,
94 default_mode: None,
95 env: Default::default(),
96 favorite_models: Vec::new(),
97 default_config_options: Default::default(),
98 favorite_config_option_values: Default::default(),
99 });
100
101 match settings {
102 settings::CustomAgentServerSettings::Custom {
103 favorite_config_option_values,
104 ..
105 }
106 | settings::CustomAgentServerSettings::Extension {
107 favorite_config_option_values,
108 ..
109 }
110 | settings::CustomAgentServerSettings::Registry {
111 favorite_config_option_values,
112 ..
113 } => {
114 let entry = favorite_config_option_values
115 .entry(config_id.clone())
116 .or_insert_with(Vec::new);
117
118 if should_be_favorite {
119 if !entry.iter().any(|v| v == &value_id) {
120 entry.push(value_id.clone());
121 }
122 } else {
123 entry.retain(|v| v != &value_id);
124 if entry.is_empty() {
125 favorite_config_option_values.remove(&config_id);
126 }
127 }
128 }
129 }
130 });
131 }
132
133 fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
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 .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 .get(self.name().as_ref())
164 .cloned()
165 });
166
167 settings
168 .as_ref()
169 .and_then(|s| s.default_model().map(acp::ModelId::new))
170 }
171
172 fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
173 let name = self.name();
174 update_settings_file(fs, cx, move |settings, _| {
175 let settings = settings
176 .agent_servers
177 .get_or_insert_default()
178 .entry(name.to_string())
179 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
180 default_model: None,
181 default_mode: None,
182 env: Default::default(),
183 favorite_models: Vec::new(),
184 default_config_options: Default::default(),
185 favorite_config_option_values: Default::default(),
186 });
187
188 match settings {
189 settings::CustomAgentServerSettings::Custom { default_model, .. }
190 | settings::CustomAgentServerSettings::Extension { default_model, .. }
191 | settings::CustomAgentServerSettings::Registry { default_model, .. } => {
192 *default_model = model_id.map(|m| m.to_string());
193 }
194 }
195 });
196 }
197
198 fn favorite_model_ids(&self, cx: &mut App) -> HashSet<acp::ModelId> {
199 let settings = cx.read_global(|settings: &SettingsStore, _| {
200 settings
201 .get::<AllAgentServersSettings>(None)
202 .get(self.name().as_ref())
203 .cloned()
204 });
205
206 settings
207 .as_ref()
208 .map(|s| {
209 s.favorite_models()
210 .iter()
211 .map(|id| acp::ModelId::new(id.clone()))
212 .collect()
213 })
214 .unwrap_or_default()
215 }
216
217 fn toggle_favorite_model(
218 &self,
219 model_id: acp::ModelId,
220 should_be_favorite: bool,
221 fs: Arc<dyn Fs>,
222 cx: &App,
223 ) {
224 let name = self.name();
225 update_settings_file(fs, cx, move |settings, _| {
226 let settings = settings
227 .agent_servers
228 .get_or_insert_default()
229 .entry(name.to_string())
230 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
231 default_model: None,
232 default_mode: None,
233 env: Default::default(),
234 favorite_models: Vec::new(),
235 default_config_options: Default::default(),
236 favorite_config_option_values: Default::default(),
237 });
238
239 let favorite_models = match settings {
240 settings::CustomAgentServerSettings::Custom {
241 favorite_models, ..
242 }
243 | settings::CustomAgentServerSettings::Extension {
244 favorite_models, ..
245 }
246 | settings::CustomAgentServerSettings::Registry {
247 favorite_models, ..
248 } => favorite_models,
249 };
250
251 let model_id_str = model_id.to_string();
252 if should_be_favorite {
253 if !favorite_models.contains(&model_id_str) {
254 favorite_models.push(model_id_str);
255 }
256 } else {
257 favorite_models.retain(|id| id != &model_id_str);
258 }
259 });
260 }
261
262 fn default_config_option(&self, config_id: &str, cx: &App) -> Option<String> {
263 let settings = cx.read_global(|settings: &SettingsStore, _| {
264 settings
265 .get::<AllAgentServersSettings>(None)
266 .get(self.name().as_ref())
267 .cloned()
268 });
269
270 settings
271 .as_ref()
272 .and_then(|s| s.default_config_option(config_id).map(|s| s.to_string()))
273 }
274
275 fn set_default_config_option(
276 &self,
277 config_id: &str,
278 value_id: Option<&str>,
279 fs: Arc<dyn Fs>,
280 cx: &mut App,
281 ) {
282 let name = self.name();
283 let config_id = config_id.to_string();
284 let value_id = value_id.map(|s| s.to_string());
285 update_settings_file(fs, cx, move |settings, _| {
286 let settings = settings
287 .agent_servers
288 .get_or_insert_default()
289 .entry(name.to_string())
290 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
291 default_model: None,
292 default_mode: None,
293 env: Default::default(),
294 favorite_models: Vec::new(),
295 default_config_options: Default::default(),
296 favorite_config_option_values: Default::default(),
297 });
298
299 match settings {
300 settings::CustomAgentServerSettings::Custom {
301 default_config_options,
302 ..
303 }
304 | settings::CustomAgentServerSettings::Extension {
305 default_config_options,
306 ..
307 }
308 | settings::CustomAgentServerSettings::Registry {
309 default_config_options,
310 ..
311 } => {
312 if let Some(value) = value_id.clone() {
313 default_config_options.insert(config_id.clone(), value);
314 } else {
315 default_config_options.remove(&config_id);
316 }
317 }
318 }
319 });
320 }
321
322 fn connect(
323 &self,
324 delegate: AgentServerDelegate,
325 cx: &mut App,
326 ) -> Task<Result<Rc<dyn AgentConnection>>> {
327 let name = self.name();
328 let display_name = delegate
329 .store
330 .read(cx)
331 .agent_display_name(&ExternalAgentServerName(name.clone()))
332 .unwrap_or_else(|| name.clone());
333 let default_mode = self.default_mode(cx);
334 let default_model = self.default_model(cx);
335 let is_previous_built_in =
336 matches!(name.as_ref(), CLAUDE_AGENT_NAME | CODEX_NAME | GEMINI_NAME);
337 let (default_config_options, is_registry_agent) =
338 cx.read_global(|settings: &SettingsStore, _| {
339 let agent_settings = settings
340 .get::<AllAgentServersSettings>(None)
341 .get(self.name().as_ref());
342
343 let is_registry = agent_settings
344 .map(|s| {
345 matches!(
346 s,
347 project::agent_server_store::CustomAgentServerSettings::Registry { .. }
348 )
349 })
350 .unwrap_or(false);
351
352 let config_options = agent_settings
353 .map(|s| match s {
354 project::agent_server_store::CustomAgentServerSettings::Custom {
355 default_config_options,
356 ..
357 }
358 | project::agent_server_store::CustomAgentServerSettings::Extension {
359 default_config_options,
360 ..
361 }
362 | project::agent_server_store::CustomAgentServerSettings::Registry {
363 default_config_options,
364 ..
365 } => default_config_options.clone(),
366 })
367 .unwrap_or_default();
368
369 (config_options, is_registry)
370 });
371
372 // Intermediate step to allow for previous built-ins to also be triggered if they aren't in settings yet.
373 let is_registry_agent = is_registry_agent || is_previous_built_in;
374
375 if is_registry_agent {
376 if let Some(registry_store) = project::AgentRegistryStore::try_global(cx) {
377 registry_store.update(cx, |store, cx| store.refresh_if_stale(cx));
378 }
379 }
380
381 let mut extra_env = load_proxy_env(cx);
382 if delegate.store.read(cx).no_browser() {
383 extra_env.insert("NO_BROWSER".to_owned(), "1".to_owned());
384 }
385 if is_registry_agent {
386 match name.as_ref() {
387 CLAUDE_AGENT_NAME => {
388 extra_env.insert("ANTHROPIC_API_KEY".into(), "".into());
389 }
390 CODEX_NAME => {
391 if let Ok(api_key) = std::env::var("CODEX_API_KEY") {
392 extra_env.insert("CODEX_API_KEY".into(), api_key);
393 }
394 if let Ok(api_key) = std::env::var("OPEN_AI_API_KEY") {
395 extra_env.insert("OPEN_AI_API_KEY".into(), api_key);
396 }
397 }
398 GEMINI_NAME => {
399 extra_env.insert("SURFACE".to_owned(), "zed".to_owned());
400 }
401 _ => {}
402 }
403 }
404 let store = delegate.store.downgrade();
405 cx.spawn(async move |cx| {
406 if is_registry_agent && name.as_ref() == GEMINI_NAME {
407 if let Some(api_key) = cx.update(api_key_for_gemini_cli).await.ok() {
408 extra_env.insert("GEMINI_API_KEY".into(), api_key);
409 }
410 }
411 let command = store
412 .update(cx, |store, cx| {
413 let agent = store
414 .get_external_agent(&ExternalAgentServerName(name.clone()))
415 .with_context(|| {
416 format!("Custom agent server `{}` is not registered", name)
417 })?;
418 anyhow::Ok(agent.get_command(
419 extra_env,
420 delegate.status_tx,
421 delegate.new_version_available,
422 &mut cx.to_async(),
423 ))
424 })??
425 .await?;
426 let connection = crate::acp::connect(
427 name,
428 display_name,
429 command,
430 default_mode,
431 default_model,
432 default_config_options,
433 cx,
434 )
435 .await?;
436 Ok(connection)
437 })
438 }
439
440 fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
441 self
442 }
443}
444
445fn api_key_for_gemini_cli(cx: &mut App) -> Task<Result<String>> {
446 let env_var = EnvVar::new("GEMINI_API_KEY".into()).or(EnvVar::new("GOOGLE_AI_API_KEY".into()));
447 if let Some(key) = env_var.value {
448 return Task::ready(Ok(key));
449 }
450 let credentials_provider = <dyn CredentialsProvider>::global(cx);
451 let api_url = google_ai::API_URL.to_string();
452 cx.spawn(async move |cx| {
453 Ok(
454 ApiKey::load_from_system_keychain(&api_url, credentials_provider.as_ref(), cx)
455 .await?
456 .key()
457 .to_string(),
458 )
459 })
460}