tenant.openagents/omega

1use crate::{AgentServer, AgentServerDelegate, load_proxy_env};
2use acp_thread::AgentConnection;
3use agent_client_protocol::schema::v1 as acp;
4use anyhow::{Context as _, Result};
5use collections::HashSet;
6use fs::Fs;
7use gpui::{App, AppContext as _, Entity, Task};
8use language_model::{ApiKey, EnvVar};
9use project::{
10    Project,
11    agent_server_store::{AgentId, AllAgentServersSettings},
12};
13use settings::{AgentConfigOptionValue, SettingsStore, update_settings_file};
14use std::{rc::Rc, sync::Arc};
15use ui::IconName;
16
17pub const GEMINI_ID: &str = "gemini";
18pub const CLAUDE_AGENT_ID: &str = "claude-acp";
19pub const CODEX_ID: &str = "codex-acp";
20pub const CURSOR_ID: &str = "cursor";
21
22/// A generic agent server implementation for custom user-defined agents
23pub struct CustomAgentServer {
24    agent_id: AgentId,
25}
26
27impl CustomAgentServer {
28    pub fn new(agent_id: AgentId) -> Self {
29        Self { agent_id }
30    }
31}
32
33impl AgentServer for CustomAgentServer {
34    fn agent_id(&self) -> AgentId {
35        self.agent_id.clone()
36    }
37
38    fn logo(&self) -> IconName {
39        IconName::Terminal
40    }
41
42    fn default_mode(&self, cx: &App) -> Option<acp::SessionModeId> {
43        let settings = cx.read_global(|settings: &SettingsStore, _| {
44            settings
45                .get::<AllAgentServersSettings>(None)
46                .get(self.agent_id().0.as_ref())
47                .cloned()
48        });
49
50        settings
51            .as_ref()
52            .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
53    }
54
55    fn favorite_config_option_value_ids(
56        &self,
57        config_id: &acp::SessionConfigId,
58        cx: &mut App,
59    ) -> HashSet<acp::SessionConfigValueId> {
60        let settings = cx.read_global(|settings: &SettingsStore, _| {
61            settings
62                .get::<AllAgentServersSettings>(None)
63                .get(self.agent_id().0.as_ref())
64                .cloned()
65        });
66
67        settings
68            .as_ref()
69            .and_then(|s| s.favorite_config_option_values(config_id.0.as_ref()))
70            .map(|values| {
71                values
72                    .iter()
73                    .cloned()
74                    .map(acp::SessionConfigValueId::new)
75                    .collect()
76            })
77            .unwrap_or_default()
78    }
79
80    fn toggle_favorite_config_option_value(
81        &self,
82        config_id: acp::SessionConfigId,
83        value_id: acp::SessionConfigValueId,
84        should_be_favorite: bool,
85        fs: Arc<dyn Fs>,
86        cx: &App,
87    ) {
88        let agent_id = self.agent_id();
89        let config_id = config_id.to_string();
90        let value_id = value_id.to_string();
91
92        update_settings_file(fs, cx, move |settings, _cx| {
93            let settings = settings
94                .agent_servers
95                .get_or_insert_default()
96                .entry(agent_id.0.to_string())
97                .or_insert_with(default_settings_for_agent);
98
99            match settings {
100                settings::CustomAgentServerSettings::Custom {
101                    favorite_config_option_values,
102                    ..
103                }
104                | settings::CustomAgentServerSettings::Registry {
105                    favorite_config_option_values,
106                    ..
107                } => {
108                    let entry = favorite_config_option_values
109                        .entry(config_id.clone())
110                        .or_insert_with(Vec::new);
111
112                    if should_be_favorite {
113                        if !entry.iter().any(|v| v == &value_id) {
114                            entry.push(value_id.clone());
115                        }
116                    } else {
117                        entry.retain(|v| v != &value_id);
118                        if entry.is_empty() {
119                            favorite_config_option_values.remove(&config_id);
120                        }
121                    }
122                }
123            }
124        });
125    }
126
127    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
128        let agent_id = self.agent_id();
129        update_settings_file(fs, cx, move |settings, _cx| {
130            let settings = settings
131                .agent_servers
132                .get_or_insert_default()
133                .entry(agent_id.0.to_string())
134                .or_insert_with(default_settings_for_agent);
135
136            match settings {
137                settings::CustomAgentServerSettings::Custom { default_mode, .. }
138                | settings::CustomAgentServerSettings::Registry { default_mode, .. } => {
139                    *default_mode = mode_id.map(|m| m.to_string());
140                }
141            }
142        });
143    }
144
145    fn default_config_option(&self, config_id: &str, cx: &App) -> Option<AgentConfigOptionValue> {
146        let settings = cx.read_global(|settings: &SettingsStore, _| {
147            settings
148                .get::<AllAgentServersSettings>(None)
149                .get(self.agent_id().as_ref())
150                .cloned()
151        });
152
153        settings
154            .as_ref()
155            .and_then(|s| s.default_config_option(config_id).cloned())
156    }
157
158    fn set_default_config_option(
159        &self,
160        config_id: &str,
161        value: Option<AgentConfigOptionValue>,
162        fs: Arc<dyn Fs>,
163        cx: &mut App,
164    ) {
165        let agent_id = self.agent_id();
166        let config_id = config_id.to_string();
167        update_settings_file(fs, cx, move |settings, _cx| {
168            let settings = settings
169                .agent_servers
170                .get_or_insert_default()
171                .entry(agent_id.0.to_string())
172                .or_insert_with(default_settings_for_agent);
173
174            match settings {
175                settings::CustomAgentServerSettings::Custom {
176                    default_config_options,
177                    ..
178                }
179                | settings::CustomAgentServerSettings::Registry {
180                    default_config_options,
181                    ..
182                } => {
183                    if let Some(value) = value {
184                        default_config_options.insert(config_id.clone(), value);
185                    } else {
186                        default_config_options.remove(&config_id);
187                    }
188                }
189            }
190        });
191    }
192
193    fn connect(
194        &self,
195        delegate: AgentServerDelegate,
196        project: Entity<Project>,
197        cx: &mut App,
198    ) -> Task<Result<Rc<dyn AgentConnection>>> {
199        let agent_id = self.agent_id();
200        let default_mode = self.default_mode(cx);
201        let is_registry_agent = is_registry_agent(agent_id.clone(), cx);
202        let default_config_options = cx.read_global(|settings: &SettingsStore, _| {
203            settings
204                .get::<AllAgentServersSettings>(None)
205                .get(self.agent_id().as_ref())
206                .map(|s| match s {
207                    project::agent_server_store::CustomAgentServerSettings::Custom {
208                        default_config_options,
209                        ..
210                    }
211                    | project::agent_server_store::CustomAgentServerSettings::Registry {
212                        default_config_options,
213                        ..
214                    } => default_config_options.clone(),
215                })
216                .unwrap_or_default()
217        });
218
219        if is_registry_agent {
220            if let Some(registry_store) = project::AgentRegistryStore::try_global(cx) {
221                registry_store.update(cx, |store, cx| store.refresh_if_stale(cx));
222            }
223        }
224
225        let mut extra_env = load_proxy_env(cx);
226        if delegate.store.read(cx).no_browser() {
227            extra_env.insert("NO_BROWSER".to_owned(), "1".to_owned());
228        }
229        if is_registry_agent {
230            match agent_id.as_ref() {
231                CLAUDE_AGENT_ID => {
232                    extra_env.insert("ANTHROPIC_API_KEY".into(), "".into());
233                }
234                CODEX_ID => {
235                    if let Ok(api_key) = std::env::var("CODEX_API_KEY") {
236                        extra_env.insert("CODEX_API_KEY".into(), api_key);
237                    }
238                    if let Ok(api_key) = std::env::var("OPEN_AI_API_KEY") {
239                        extra_env.insert("OPEN_AI_API_KEY".into(), api_key);
240                    }
241                }
242                GEMINI_ID => {
243                    extra_env.insert("SURFACE".to_owned(), "zed".to_owned());
244                }
245                _ => {}
246            }
247        }
248        let store = delegate.store.downgrade();
249        cx.spawn(async move |cx| {
250            if is_registry_agent && agent_id.as_ref() == GEMINI_ID {
251                if let Some(api_key) = cx.update(api_key_for_gemini_cli).await.ok() {
252                    extra_env.insert("GEMINI_API_KEY".into(), api_key);
253                }
254            }
255            let command = store
256                .update(cx, |store, cx| {
257                    let agent = store.get_external_agent(&agent_id).with_context(|| {
258                        format!("Custom agent server `{}` is not registered", agent_id)
259                    })?;
260                    if let Some(new_version_available_tx) = delegate.new_version_available {
261                        agent.set_new_version_available_tx(new_version_available_tx);
262                    }
263                    if let Some(loading_status_tx) = delegate.loading_status {
264                        agent.set_loading_status_tx(loading_status_tx);
265                    }
266                    anyhow::Ok(agent.get_command(vec![], extra_env, &mut cx.to_async()))
267                })??
268                .await?;
269            let connection = crate::acp::connect(
270                agent_id,
271                project,
272                command,
273                store.clone(),
274                default_mode,
275                default_config_options,
276                cx,
277            )
278            .await?;
279            Ok(connection)
280        })
281    }
282
283    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
284        self
285    }
286}
287
288fn api_key_for_gemini_cli(cx: &mut App) -> Task<Result<String>> {
289    let env_var = EnvVar::new("GEMINI_API_KEY".into()).or(EnvVar::new("GOOGLE_AI_API_KEY".into()));
290    if let Some(key) = env_var.value {
291        return Task::ready(Ok(key));
292    }
293    let credentials_provider = zed_credentials_provider::global(cx);
294    let api_url = google_ai::API_URL.to_string();
295    cx.spawn(async move |cx| {
296        Ok(
297            ApiKey::load_from_system_keychain(&api_url, credentials_provider.as_ref(), cx)
298                .await?
299                .key()
300                .to_string(),
301        )
302    })
303}
304
305fn is_registry_agent(agent_id: impl Into<AgentId>, cx: &App) -> bool {
306    let agent_id = agent_id.into();
307    let is_in_registry = project::AgentRegistryStore::try_global(cx)
308        .map(|store| store.read(cx).agent(&agent_id).is_some())
309        .unwrap_or(false);
310    let is_settings_registry = cx.read_global(|settings: &SettingsStore, _| {
311        settings
312            .get::<AllAgentServersSettings>(None)
313            .get(agent_id.as_ref())
314            .is_some_and(|s| {
315                matches!(
316                    s,
317                    project::agent_server_store::CustomAgentServerSettings::Registry { .. }
318                )
319            })
320    });
321    is_in_registry || is_settings_registry
322}
323
324fn default_settings_for_agent() -> settings::CustomAgentServerSettings {
325    settings::CustomAgentServerSettings::Registry {
326        default_mode: None,
327        env: Default::default(),
328        default_config_options: Default::default(),
329        favorite_config_option_values: Default::default(),
330    }
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336    use collections::HashMap;
337    use gpui::TestAppContext;
338    use project::agent_registry_store::{
339        AgentRegistryStore, RegistryAgent, RegistryAgentMetadata, RegistryNpxAgent,
340    };
341    use settings::Settings as _;
342    use ui::SharedString;
343
344    fn init_test(cx: &mut TestAppContext) {
345        cx.update(|cx| {
346            let settings_store = SettingsStore::test(cx);
347            cx.set_global(settings_store);
348        });
349    }
350
351    fn init_registry_with_agents(cx: &mut TestAppContext, agent_ids: &[&str]) {
352        let agents: Vec<RegistryAgent> = agent_ids
353            .iter()
354            .map(|id| {
355                let id = SharedString::from(id.to_string());
356                RegistryAgent::Npx(RegistryNpxAgent {
357                    metadata: RegistryAgentMetadata {
358                        id: AgentId::new(id.clone()),
359                        name: id.clone(),
360                        description: SharedString::from(""),
361                        version: SharedString::from("1.0.0"),
362                        repository: None,
363                        website: None,
364                        icon_path: None,
365                    },
366                    package: id,
367                    args: Vec::new(),
368                    env: HashMap::default(),
369                })
370            })
371            .collect();
372        cx.update(|cx| {
373            AgentRegistryStore::init_test_global(cx, agents);
374        });
375    }
376
377    fn set_agent_server_settings(
378        cx: &mut TestAppContext,
379        entries: Vec<(&str, settings::CustomAgentServerSettings)>,
380    ) {
381        cx.update(|cx| {
382            AllAgentServersSettings::override_global(
383                project::agent_server_store::AllAgentServersSettings(
384                    entries
385                        .into_iter()
386                        .map(|(name, settings)| (name.to_string(), settings.into()))
387                        .collect(),
388                ),
389                cx,
390            );
391        });
392    }
393
394    #[gpui::test]
395    fn test_unknown_agent_is_not_registry(cx: &mut TestAppContext) {
396        init_test(cx);
397        cx.update(|cx| {
398            assert!(!is_registry_agent("my-custom-agent", cx));
399        });
400    }
401
402    #[gpui::test]
403    fn test_agent_in_registry_store_is_registry(cx: &mut TestAppContext) {
404        init_test(cx);
405        init_registry_with_agents(cx, &["some-new-registry-agent"]);
406        cx.update(|cx| {
407            assert!(is_registry_agent("some-new-registry-agent", cx));
408            assert!(!is_registry_agent("not-in-registry", cx));
409        });
410    }
411
412    #[gpui::test]
413    fn test_agent_with_registry_settings_type_is_registry(cx: &mut TestAppContext) {
414        init_test(cx);
415        set_agent_server_settings(
416            cx,
417            vec![(
418                "agent-from-settings",
419                settings::CustomAgentServerSettings::Registry {
420                    env: HashMap::default(),
421                    default_mode: None,
422                    default_config_options: HashMap::default(),
423                    favorite_config_option_values: HashMap::default(),
424                },
425            )],
426        );
427        cx.update(|cx| {
428            assert!(is_registry_agent("agent-from-settings", cx));
429        });
430    }
431}
432