tenant.openagents/omega

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use anyhow::{Context as _, Result, anyhow, bail};
6use collections::HashMap;
7use fs::Fs;
8use futures::{AsyncReadExt, future::join_all};
9use gpui::{
10    App, AppContext as _, BackgroundExecutor, Context, Entity, FutureExt as _, Global,
11    SharedString, Task, TaskExt,
12};
13use http_client::{AsyncBody, HttpClient, StatusCode};
14use serde::Deserialize;
15use settings::Settings as _;
16use util::ResultExt;
17
18use crate::{AgentId, DisableAiSettings};
19
20const REGISTRY_URL: &str = "https://cdn.agentclientprotocol.com/registry/v1/latest/registry.json";
21const REFRESH_THROTTLE_DURATION: Duration = Duration::from_secs(60 * 60);
22// Bound the full request lifecycle, including response body reads; the shared
23// HTTP client only has a connect timeout.
24const REGISTRY_FETCH_TIMEOUT: Duration = Duration::from_secs(30);
25const REGISTRY_ICON_FETCH_TIMEOUT: Duration = Duration::from_secs(10);
26
27#[derive(Clone, Debug)]
28pub struct RegistryAgentMetadata {
29    pub id: AgentId,
30    pub name: SharedString,
31    pub description: SharedString,
32    pub version: SharedString,
33    pub repository: Option<SharedString>,
34    pub website: Option<SharedString>,
35    pub icon_path: Option<SharedString>,
36}
37
38#[derive(Clone, Debug)]
39pub struct RegistryBinaryAgent {
40    pub metadata: RegistryAgentMetadata,
41    pub targets: HashMap<String, RegistryTargetConfig>,
42    pub supports_current_platform: bool,
43}
44
45#[derive(Clone, Debug)]
46pub struct RegistryNpxAgent {
47    pub metadata: RegistryAgentMetadata,
48    pub package: SharedString,
49    pub args: Vec<String>,
50    pub env: HashMap<String, String>,
51}
52
53#[derive(Clone, Debug)]
54pub enum RegistryAgent {
55    Binary(RegistryBinaryAgent),
56    Npx(RegistryNpxAgent),
57}
58
59impl RegistryAgent {
60    pub fn metadata(&self) -> &RegistryAgentMetadata {
61        match self {
62            RegistryAgent::Binary(agent) => &agent.metadata,
63            RegistryAgent::Npx(agent) => &agent.metadata,
64        }
65    }
66
67    pub fn id(&self) -> &AgentId {
68        &self.metadata().id
69    }
70
71    pub fn name(&self) -> &SharedString {
72        &self.metadata().name
73    }
74
75    pub fn description(&self) -> &SharedString {
76        &self.metadata().description
77    }
78
79    pub fn version(&self) -> &SharedString {
80        &self.metadata().version
81    }
82
83    pub fn repository(&self) -> Option<&SharedString> {
84        self.metadata().repository.as_ref()
85    }
86
87    pub fn website(&self) -> Option<&SharedString> {
88        self.metadata().website.as_ref()
89    }
90
91    pub fn icon_path(&self) -> Option<&SharedString> {
92        self.metadata().icon_path.as_ref()
93    }
94
95    pub fn supports_current_platform(&self) -> bool {
96        match self {
97            RegistryAgent::Binary(agent) => agent.supports_current_platform,
98            RegistryAgent::Npx(_) => true,
99        }
100    }
101}
102
103#[derive(Clone, Debug)]
104pub struct RegistryTargetConfig {
105    pub archive: String,
106    pub cmd: String,
107    pub args: Vec<String>,
108    pub sha256: Option<String>,
109    pub env: HashMap<String, String>,
110}
111
112struct GlobalAgentRegistryStore(Entity<AgentRegistryStore>);
113
114impl Global for GlobalAgentRegistryStore {}
115
116pub struct AgentRegistryStore {
117    fs: Arc<dyn Fs>,
118    http_client: Arc<dyn HttpClient>,
119    agents: Vec<RegistryAgent>,
120    is_fetching: bool,
121    fetch_error: Option<SharedString>,
122    pending_refresh: Option<Task<()>>,
123    last_refresh: Option<Instant>,
124}
125
126impl AgentRegistryStore {
127    /// Initialize the global AgentRegistryStore.
128    ///
129    /// This loads the cached registry from disk. If the cache is empty but there
130    /// are registry agents configured in settings, it will trigger a network fetch.
131    /// Otherwise, call `refresh()` explicitly when you need fresh data
132    /// (e.g., when opening the Agent Registry page).
133    pub fn init_global(
134        cx: &mut App,
135        fs: Arc<dyn Fs>,
136        http_client: Arc<dyn HttpClient>,
137    ) -> Entity<Self> {
138        if let Some(store) = Self::try_global(cx) {
139            return store;
140        }
141
142        let store = cx.new(|cx| Self::new(fs, http_client, cx));
143        cx.set_global(GlobalAgentRegistryStore(store.clone()));
144
145        store.update(cx, |store, cx| {
146            if store.agents.is_empty() {
147                store.refresh(cx);
148            }
149        });
150
151        store
152    }
153
154    pub fn global(cx: &App) -> Entity<Self> {
155        cx.global::<GlobalAgentRegistryStore>().0.clone()
156    }
157
158    pub fn try_global(cx: &App) -> Option<Entity<Self>> {
159        cx.try_global::<GlobalAgentRegistryStore>()
160            .map(|store| store.0.clone())
161    }
162
163    #[cfg(any(test, feature = "test-support"))]
164    pub fn init_test_global(cx: &mut App, agents: Vec<RegistryAgent>) -> Entity<Self> {
165        let fs: Arc<dyn Fs> = fs::FakeFs::new(cx.background_executor().clone());
166        let store = cx.new(|_cx| Self {
167            fs,
168            http_client: http_client::FakeHttpClient::with_404_response(),
169            agents,
170            is_fetching: false,
171            fetch_error: None,
172            pending_refresh: None,
173            last_refresh: None,
174        });
175        cx.set_global(GlobalAgentRegistryStore(store.clone()));
176        store
177    }
178
179    #[cfg(any(test, feature = "test-support"))]
180    pub fn set_agents(&mut self, agents: Vec<RegistryAgent>, cx: &mut Context<Self>) {
181        self.agents = agents;
182        cx.notify();
183    }
184
185    pub fn agents(&self) -> &[RegistryAgent] {
186        &self.agents
187    }
188
189    pub fn agent(&self, id: &AgentId) -> Option<&RegistryAgent> {
190        self.agents.iter().find(|agent| agent.id() == id)
191    }
192
193    pub fn is_fetching(&self) -> bool {
194        self.is_fetching
195    }
196
197    pub fn fetch_error(&self) -> Option<SharedString> {
198        self.fetch_error.clone()
199    }
200
201    /// Refresh the registry from the network.
202    ///
203    /// This will fetch the latest registry data and update the cache.
204    pub fn refresh(&mut self, cx: &mut Context<Self>) {
205        if self.pending_refresh.is_some() {
206            return;
207        }
208
209        if DisableAiSettings::get_global(cx).disable_ai {
210            return;
211        }
212
213        self.is_fetching = true;
214        self.fetch_error = None;
215        self.last_refresh = Some(Instant::now());
216        cx.notify();
217
218        let fs = self.fs.clone();
219        let http_client = self.http_client.clone();
220        let executor = cx.background_executor().clone();
221
222        self.pending_refresh = Some(cx.spawn(async move |this, cx| {
223            let result = match fetch_registry_index(http_client.clone(), &executor).await {
224                Ok(data) => {
225                    build_registry_agents(
226                        fs.clone(),
227                        http_client,
228                        data.index,
229                        data.raw_body,
230                        true,
231                        &executor,
232                    )
233                    .await
234                }
235                Err(error) => {
236                    log::error!("AgentRegistryStore::refresh: fetch failed: {error:#}");
237                    Err(error)
238                }
239            };
240
241            this.update(cx, |this, cx| {
242                this.pending_refresh = None;
243                this.is_fetching = false;
244                match result {
245                    Ok(agents) => {
246                        this.agents = agents;
247                        this.fetch_error = None;
248                    }
249                    Err(error) => {
250                        this.fetch_error = Some(SharedString::from(format!("{error:#}")));
251                    }
252                }
253                cx.notify();
254            })
255            .ok();
256        }));
257    }
258
259    /// Refresh the registry if it hasn't been refreshed recently.
260    ///
261    /// This is useful to call when using a registry-based agent to check for
262    /// updates without making too many network requests. The refresh is
263    /// throttled to at most once per hour.
264    pub fn refresh_if_stale(&mut self, cx: &mut Context<Self>) {
265        let should_refresh = self
266            .last_refresh
267            .map(|last| last.elapsed() >= REFRESH_THROTTLE_DURATION)
268            .unwrap_or(true);
269
270        if should_refresh {
271            self.refresh(cx);
272        }
273    }
274
275    fn new(fs: Arc<dyn Fs>, http_client: Arc<dyn HttpClient>, cx: &mut Context<Self>) -> Self {
276        let mut store = Self {
277            fs: fs.clone(),
278            http_client,
279            agents: Vec::new(),
280            is_fetching: false,
281            fetch_error: None,
282            pending_refresh: None,
283            last_refresh: None,
284        };
285
286        store.load_cached_registry(fs, store.http_client.clone(), cx);
287
288        store
289    }
290
291    fn load_cached_registry(
292        &mut self,
293        fs: Arc<dyn Fs>,
294        http_client: Arc<dyn HttpClient>,
295        cx: &mut Context<Self>,
296    ) {
297        if DisableAiSettings::get_global(cx).disable_ai {
298            return;
299        }
300
301        cx.spawn(async move |this, cx| -> Result<()> {
302            let cache_path = registry_cache_path();
303            if !fs.is_file(&cache_path).await {
304                return Ok(());
305            }
306
307            let bytes = fs
308                .load_bytes(&cache_path)
309                .await
310                .context("reading cached registry")?;
311            let index: RegistryIndex =
312                serde_json::from_slice(&bytes).context("parsing cached registry")?;
313
314            let executor = cx.background_executor().clone();
315            let agents =
316                build_registry_agents(fs, http_client, index, bytes, false, &executor).await?;
317
318            this.update(cx, |this, cx| {
319                this.agents = agents;
320                cx.notify();
321            })?;
322
323            Ok(())
324        })
325        .detach_and_log_err(cx);
326    }
327}
328
329struct RegistryFetchResult {
330    index: RegistryIndex,
331    raw_body: Vec<u8>,
332}
333
334async fn fetch_registry_index(
335    http_client: Arc<dyn HttpClient>,
336    executor: &BackgroundExecutor,
337) -> Result<RegistryFetchResult> {
338    let (status, body) =
339        fetch_url_body(http_client, REGISTRY_URL, REGISTRY_FETCH_TIMEOUT, executor)
340            .await
341            .context("fetching ACP registry")?;
342
343    if status.is_client_error() {
344        let text = String::from_utf8_lossy(body.as_slice());
345        bail!(
346            "registry status error {}, response: {text:?}",
347            status.as_u16()
348        );
349    }
350
351    let index: RegistryIndex = serde_json::from_slice(&body).context("parsing ACP registry")?;
352    Ok(RegistryFetchResult {
353        index,
354        raw_body: body,
355    })
356}
357
358async fn build_registry_agents(
359    fs: Arc<dyn Fs>,
360    http_client: Arc<dyn HttpClient>,
361    index: RegistryIndex,
362    raw_body: Vec<u8>,
363    update_cache: bool,
364    executor: &BackgroundExecutor,
365) -> Result<Vec<RegistryAgent>> {
366    let cache_dir = registry_cache_dir();
367    fs.create_dir(&cache_dir).await?;
368
369    let cache_path = cache_dir.join("registry.json");
370    if update_cache {
371        fs.write(&cache_path, &raw_body).await?;
372    }
373
374    let icons_dir = cache_dir.join("icons");
375    if update_cache {
376        fs.create_dir(&icons_dir).await?;
377    }
378
379    let current_platform = current_platform_key();
380    let icon_paths = resolve_icon_paths(
381        &index.agents,
382        &icons_dir,
383        update_cache,
384        fs.clone(),
385        http_client.clone(),
386        executor,
387    )
388    .await;
389
390    let mut agents = Vec::new();
391    for (entry, icon_path) in index.agents.into_iter().zip(icon_paths) {
392        let metadata = RegistryAgentMetadata {
393            id: AgentId::new(entry.id),
394            name: entry.name.into(),
395            description: entry.description.into(),
396            version: entry.version.into(),
397            repository: entry.repository.map(Into::into),
398            website: entry.website.map(Into::into),
399            icon_path,
400        };
401
402        let binary_agent = entry.distribution.binary.as_ref().and_then(|binary| {
403            if binary.is_empty() {
404                return None;
405            }
406
407            let mut targets = HashMap::default();
408            for (platform, target) in binary.iter() {
409                targets.insert(
410                    platform.clone(),
411                    RegistryTargetConfig {
412                        archive: target.archive.clone(),
413                        cmd: target.cmd.clone(),
414                        args: target.args.clone(),
415                        sha256: target.sha256.clone(),
416                        env: target.env.clone(),
417                    },
418                );
419            }
420
421            let supports_current_platform = current_platform
422                .as_ref()
423                .is_some_and(|platform| targets.contains_key(*platform));
424
425            Some(RegistryBinaryAgent {
426                metadata: metadata.clone(),
427                targets,
428                supports_current_platform,
429            })
430        });
431
432        let npx_agent = entry.distribution.npx.as_ref().map(|npx| RegistryNpxAgent {
433            metadata: metadata.clone(),
434            package: npx.package.clone().into(),
435            args: npx.args.clone(),
436            env: npx.env.clone(),
437        });
438
439        let agent = match (binary_agent, npx_agent) {
440            (Some(binary_agent), Some(npx_agent)) => {
441                if binary_agent.supports_current_platform {
442                    RegistryAgent::Binary(binary_agent)
443                } else {
444                    RegistryAgent::Npx(npx_agent)
445                }
446            }
447            (Some(binary_agent), None) => RegistryAgent::Binary(binary_agent),
448            (None, Some(npx_agent)) => RegistryAgent::Npx(npx_agent),
449            (None, None) => continue,
450        };
451
452        agents.push(agent);
453    }
454
455    Ok(agents)
456}
457
458async fn resolve_icon_paths(
459    entries: &[RegistryEntry],
460    icons_dir: &Path,
461    update_cache: bool,
462    fs: Arc<dyn Fs>,
463    http_client: Arc<dyn HttpClient>,
464    executor: &BackgroundExecutor,
465) -> Vec<Option<SharedString>> {
466    join_all(entries.iter().map(|entry| {
467        let fs = fs.clone();
468        let http_client = http_client.clone();
469        async move {
470            resolve_icon_path(entry, icons_dir, update_cache, fs, http_client, executor)
471                .await
472                .log_err()
473                .flatten()
474        }
475    }))
476    .await
477}
478
479async fn resolve_icon_path(
480    entry: &RegistryEntry,
481    icons_dir: &Path,
482    update_cache: bool,
483    fs: Arc<dyn Fs>,
484    http_client: Arc<dyn HttpClient>,
485    executor: &BackgroundExecutor,
486) -> Result<Option<SharedString>> {
487    let icon_url = resolve_icon_url(entry);
488    let Some(icon_url) = icon_url else {
489        return Ok(None);
490    };
491
492    let icon_path = icons_dir.join(format!("{}.svg", entry.id));
493    if update_cache && !fs.is_file(&icon_path).await {
494        if let Err(error) = download_icon(fs.clone(), http_client, &icon_url, entry, executor).await
495        {
496            log::warn!(
497                "Failed to download ACP registry icon for {}: {error:#}",
498                entry.id
499            );
500        }
501    }
502
503    if fs.is_file(&icon_path).await {
504        Ok(Some(SharedString::from(
505            icon_path.to_string_lossy().into_owned(),
506        )))
507    } else {
508        Ok(None)
509    }
510}
511
512async fn download_icon(
513    fs: Arc<dyn Fs>,
514    http_client: Arc<dyn HttpClient>,
515    icon_url: &str,
516    entry: &RegistryEntry,
517    executor: &BackgroundExecutor,
518) -> Result<()> {
519    let (status, body) =
520        fetch_url_body(http_client, icon_url, REGISTRY_ICON_FETCH_TIMEOUT, executor)
521            .await
522            .with_context(|| format!("fetching icon for {}", entry.id))?;
523
524    if status.is_client_error() {
525        let text = String::from_utf8_lossy(body.as_slice());
526        bail!("icon status error {}, response: {text:?}", status.as_u16());
527    }
528
529    let icon_path = registry_cache_dir()
530        .join("icons")
531        .join(format!("{}.svg", entry.id));
532    fs.write(&icon_path, &body).await?;
533    Ok(())
534}
535
536async fn fetch_url_body(
537    http_client: Arc<dyn HttpClient>,
538    url: &str,
539    timeout: Duration,
540    executor: &BackgroundExecutor,
541) -> Result<(StatusCode, Vec<u8>)> {
542    async {
543        let mut response = http_client
544            .get(url, AsyncBody::default(), true)
545            .await
546            .with_context(|| format!("requesting {url}"))?;
547
548        let status = response.status();
549        let mut body = Vec::new();
550        response
551            .body_mut()
552            .read_to_end(&mut body)
553            .await
554            .with_context(|| format!("reading response from {url}"))?;
555
556        Ok((status, body))
557    }
558    .with_timeout(timeout, executor)
559    .await
560    .map_err(|_| {
561        anyhow!(
562            "timed out after {}s while fetching {url}",
563            timeout.as_secs()
564        )
565    })?
566}
567
568fn resolve_icon_url(entry: &RegistryEntry) -> Option<String> {
569    let icon = entry.icon.as_ref()?;
570    if icon.starts_with("https://") || icon.starts_with("http://") {
571        return Some(icon.to_string());
572    }
573
574    let relative_icon = icon.trim_start_matches("./");
575    Some(format!(
576        "https://raw.githubusercontent.com/agentclientprotocol/registry/main/{}/{relative_icon}",
577        entry.id
578    ))
579}
580
581fn current_platform_key() -> Option<&'static str> {
582    let os = if cfg!(target_os = "macos") {
583        "darwin"
584    } else if cfg!(target_os = "linux") {
585        "linux"
586    } else if cfg!(target_os = "windows") {
587        "windows"
588    } else {
589        return None;
590    };
591
592    let arch = if cfg!(target_arch = "aarch64") {
593        "aarch64"
594    } else if cfg!(target_arch = "x86_64") {
595        "x86_64"
596    } else {
597        return None;
598    };
599
600    Some(match os {
601        "darwin" => match arch {
602            "aarch64" => "darwin-aarch64",
603            "x86_64" => "darwin-x86_64",
604            _ => return None,
605        },
606        "linux" => match arch {
607            "aarch64" => "linux-aarch64",
608            "x86_64" => "linux-x86_64",
609            _ => return None,
610        },
611        "windows" => match arch {
612            "aarch64" => "windows-aarch64",
613            "x86_64" => "windows-x86_64",
614            _ => return None,
615        },
616        _ => return None,
617    })
618}
619
620fn registry_cache_dir() -> PathBuf {
621    paths::external_agents_dir().join("registry")
622}
623
624fn registry_cache_path() -> PathBuf {
625    registry_cache_dir().join("registry.json")
626}
627
628#[derive(Deserialize)]
629struct RegistryIndex {
630    #[serde(rename = "version")]
631    _version: String,
632    agents: Vec<RegistryEntry>,
633}
634
635#[derive(Deserialize)]
636struct RegistryEntry {
637    id: String,
638    name: String,
639    version: String,
640    description: String,
641    #[serde(default)]
642    repository: Option<String>,
643    #[serde(default)]
644    website: Option<String>,
645    #[serde(default)]
646    icon: Option<String>,
647    distribution: RegistryDistribution,
648}
649
650#[derive(Deserialize)]
651struct RegistryDistribution {
652    #[serde(default)]
653    binary: Option<HashMap<String, RegistryBinaryTarget>>,
654    #[serde(default)]
655    npx: Option<RegistryNpxDistribution>,
656}
657
658#[derive(Deserialize)]
659struct RegistryBinaryTarget {
660    archive: String,
661    cmd: String,
662    #[serde(default)]
663    args: Vec<String>,
664    #[serde(default)]
665    sha256: Option<String>,
666    #[serde(default)]
667    env: HashMap<String, String>,
668}
669
670#[derive(Deserialize)]
671struct RegistryNpxDistribution {
672    package: String,
673    #[serde(default)]
674    args: Vec<String>,
675    #[serde(default)]
676    env: HashMap<String, String>,
677}
678