use std::{ any::Any, path::{Path, PathBuf}, sync::{Arc, LazyLock}, time::Duration, }; use anyhow::{Context as _, Result, bail}; use collections::HashMap; use fs::{Fs, RemoveOptions}; use futures::StreamExt; use gpui::{ AppContext as _, AsyncApp, Context, Entity, EventEmitter, SharedString, Subscription, Task, TaskExt, }; use http_client::{HttpClient, github::AssetKind}; use node_runtime::NodeRuntime; use percent_encoding::percent_decode_str; use remote::RemoteClient; use rpc::{AnyProtoClient, TypedEnvelope, proto}; use schemars::JsonSchema; use semver::Version; use serde::{Deserialize, Serialize}; use settings::{AgentConfigOptionValue, RegisterSetting, SettingsStore, update_settings_file}; use sha2::{Digest, Sha256}; use url::Url; use util::{ResultExt as _, debug_panic}; use crate::ProjectEnvironment; use crate::agent_registry_store::{AgentRegistryStore, RegistryAgent, RegistryTargetConfig}; use crate::worktree_store::WorktreeStore; #[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema)] pub struct AgentServerCommand { #[serde(rename = "command")] pub path: PathBuf, #[serde(default)] pub args: Vec, pub env: Option>, } impl std::fmt::Debug for AgentServerCommand { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let filtered_env = self.env.as_ref().map(|env| { env.iter() .map(|(k, v)| { ( k, if util::redact::should_redact(k) { "[REDACTED]" } else { v }, ) }) .collect::>() }); f.debug_struct("AgentServerCommand") .field("path", &self.path) .field("args", &self.args) .field("env", &filtered_env) .finish() } } #[derive( Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, JsonSchema, )] #[serde(transparent)] pub struct AgentId(pub SharedString); impl AgentId { pub fn new(id: impl Into) -> Self { AgentId(id.into()) } } impl std::fmt::Display for AgentId { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) } } impl From<&'static str> for AgentId { fn from(value: &'static str) -> Self { AgentId(value.into()) } } impl From for SharedString { fn from(value: AgentId) -> Self { value.0 } } impl AsRef for AgentId { fn as_ref(&self) -> &str { &self.0 } } impl std::borrow::Borrow for AgentId { fn borrow(&self) -> &str { &self.0 } } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum ExternalAgentSource { #[default] Custom, Registry, } pub trait ExternalAgentServer { fn get_command( &mut self, extra_args: Vec, extra_env: HashMap, cx: &mut AsyncApp, ) -> Task>; fn version(&self) -> Option<&SharedString> { None } fn take_new_version_available_tx(&mut self) -> Option>> { None } fn set_new_version_available_tx(&mut self, _tx: watch::Sender>) {} fn take_loading_status_tx(&mut self) -> Option>> { None } fn set_loading_status_tx(&mut self, _tx: watch::Sender>) {} fn as_any(&self) -> &dyn Any; fn as_any_mut(&mut self) -> &mut dyn Any; } enum AgentServerStoreState { Local { node_runtime: NodeRuntime, fs: Arc, project_environment: Entity, downstream_client: Option<(u64, AnyProtoClient)>, settings: Option, http_client: Arc, _subscriptions: Vec, }, Remote { project_id: u64, upstream_client: Entity, worktree_store: Entity, }, Collab, } pub struct ExternalAgentEntry { server: Box, icon: Option, display_name: Option, pub source: ExternalAgentSource, } impl ExternalAgentEntry { pub fn new( server: Box, source: ExternalAgentSource, icon: Option, display_name: Option, ) -> Self { Self { server, icon, display_name, source, } } } pub struct AgentServerStore { state: AgentServerStoreState, pub external_agents: HashMap, } pub struct AgentServersUpdated; impl EventEmitter for AgentServerStore {} static EXTENSION_TO_REGISTRY_IDS: LazyLock> = LazyLock::new(|| { HashMap::from_iter([ ("opencode", "opencode"), ("mistral-vibe", "mistral-vibe"), ("auggie", "auggie"), ("stakpak", "stakpak"), ("codebuddy", "codebuddy-code"), ("autohand-acp", "autohand"), ("corust-agent", "corust-agent"), ("factory-droid", "factory-droid"), // Unmaintained // ("qqcode", ""), ]) }); impl AgentServerStore { pub fn migrate_agent_server_from_extensions( &mut self, id: Arc, fs: Arc, cx: &mut Context, ) { let Some(registry_id) = EXTENSION_TO_REGISTRY_IDS.get(id.as_ref()) else { return; }; update_settings_file(fs, cx, move |settings, _| { let agent_servers = settings.agent_servers.get_or_insert_default(); // Take the old settings let settings = agent_servers.remove(id.as_ref()); // If they had both installed, just remove the extension settings, leave theirregistry settings alone if agent_servers.contains_key(*registry_id) { return; } // Insert the old settings, or write new ones so it is "installed" via the registry agent_servers.insert( registry_id.to_string(), settings.unwrap_or_else(|| settings::CustomAgentServerSettings::Registry { default_mode: None, env: Default::default(), default_config_options: HashMap::default(), favorite_config_option_values: HashMap::default(), }), ); }); } pub fn agent_icon(&self, id: &AgentId) -> Option { self.external_agents .get(id) .and_then(|entry| entry.icon.clone()) } pub fn agent_source(&self, name: &AgentId) -> Option { self.external_agents.get(name).map(|entry| entry.source) } } impl AgentServerStore { pub fn agent_display_name(&self, name: &AgentId) -> Option { self.external_agents .get(name) .and_then(|entry| entry.display_name.clone()) } pub fn init_remote(session: &AnyProtoClient) { session.add_entity_message_handler(Self::handle_external_agents_updated); session.add_entity_message_handler(Self::handle_loading_status_updated); session.add_entity_message_handler(Self::handle_new_version_available); } pub fn init_headless(session: &AnyProtoClient) { session.add_entity_request_handler(Self::handle_get_agent_server_command); } fn agent_servers_settings_changed(&mut self, cx: &mut Context) { let AgentServerStoreState::Local { settings: old_settings, .. } = &mut self.state else { debug_panic!( "should not be subscribed to agent server settings changes in non-local project" ); return; }; let new_settings = cx .global::() .get::(None) .clone(); if Some(&new_settings) == old_settings.as_ref() { return; } self.reregister_agents(cx); } fn reregister_agents(&mut self, cx: &mut Context) { let AgentServerStoreState::Local { node_runtime, fs, project_environment, downstream_client, settings: old_settings, http_client, .. } = &mut self.state else { debug_panic!("Non-local projects should never attempt to reregister. This is a bug!"); return; }; let new_settings = cx .global::() .get::(None) .clone(); // If we don't have agents from the registry loaded yet, trigger a // refresh, which will cause this function to be called again let registry_store = AgentRegistryStore::try_global(cx); if new_settings.has_registry_agents() && let Some(registry) = registry_store.as_ref() { registry.update(cx, |registry, cx| registry.refresh_if_stale(cx)); } let registry_agents_by_id = registry_store .as_ref() .map(|store| { store .read(cx) .agents() .iter() .cloned() .map(|agent| (agent.id().to_string(), agent)) .collect::>() }) .unwrap_or_default(); // Drain the existing versioned agents, extracting reconnect state // from any active connection so we can preserve it or trigger a // reconnect when the version changes. let mut old_versioned_agents: HashMap< AgentId, ( SharedString, Option>>, Option>>, ), > = HashMap::default(); for (name, mut entry) in self.external_agents.drain() { if let Some(version) = entry.server.version().cloned() { let new_version_available_tx = entry.server.take_new_version_available_tx(); let loading_status_tx = entry.server.take_loading_status_tx(); if new_version_available_tx.is_some() || loading_status_tx.is_some() { old_versioned_agents .insert(name, (version, new_version_available_tx, loading_status_tx)); } } } for (name, settings) in new_settings.iter() { match settings { CustomAgentServerSettings::Custom { command, .. } => { let agent_name = AgentId(name.clone().into()); self.external_agents.insert( agent_name.clone(), ExternalAgentEntry::new( Box::new(LocalCustomAgent { command: command.clone(), project_environment: project_environment.clone(), }) as Box, ExternalAgentSource::Custom, None, None, ), ); } CustomAgentServerSettings::Registry { env, .. } => { let Some(agent) = registry_agents_by_id.get(name) else { if registry_store.is_some() { log::debug!("Registry agent '{}' not found in ACP registry", name); } continue; }; let agent_name = AgentId(name.clone().into()); match agent { RegistryAgent::Binary(agent) => { if !agent.supports_current_platform { log::warn!( "Registry agent '{}' has no compatible binary for this platform", name ); continue; } self.external_agents.insert( agent_name.clone(), ExternalAgentEntry::new( Box::new(LocalRegistryArchiveAgent { fs: fs.clone(), http_client: http_client.clone(), node_runtime: node_runtime.clone(), project_environment: project_environment.clone(), installation_dir: paths::external_agents_dir() .join("registry") .join(sanitize_path_component(name)), version: agent.metadata.version.clone(), targets: agent.targets.clone(), env: env.clone(), new_version_available_tx: None, loading_status_tx: None, }) as Box, ExternalAgentSource::Registry, agent.metadata.icon_path.clone(), Some(agent.metadata.name.clone()), ), ); } RegistryAgent::Npx(agent) => { self.external_agents.insert( agent_name.clone(), ExternalAgentEntry::new( Box::new(LocalRegistryNpxAgent { fs: fs.clone(), node_runtime: node_runtime.clone(), project_environment: project_environment.clone(), registry_id: Arc::from(name.as_str()), version: agent.metadata.version.clone(), package: agent.package.clone(), args: agent.args.clone(), distribution_env: agent.env.clone(), settings_env: env.clone(), new_version_available_tx: None, }) as Box, ExternalAgentSource::Registry, agent.metadata.icon_path.clone(), Some(agent.metadata.name.clone()), ), ); } } } } } // For each rebuilt versioned agent, compare the version. If it // changed, notify the active connection to reconnect. Otherwise, // transfer the channel to the new entry so future updates can use it. for (name, entry) in &mut self.external_agents { let Some((old_version, new_version_available_tx, loading_status_tx)) = old_versioned_agents.remove(name) else { continue; }; let Some(new_version) = entry.server.version() else { continue; }; if new_version != &old_version { if let Some(mut tx) = new_version_available_tx { tx.send(Some(new_version.to_string())).ok(); } } else { if let Some(tx) = new_version_available_tx { entry.server.set_new_version_available_tx(tx); } if let Some(tx) = loading_status_tx { entry.server.set_loading_status_tx(tx); } } } *old_settings = Some(new_settings); if let Some((project_id, downstream_client)) = downstream_client { downstream_client .send(proto::ExternalAgentsUpdated { project_id: *project_id, names: self .external_agents .keys() .map(|name| name.to_string()) .collect(), }) .log_err(); } cx.emit(AgentServersUpdated); } pub fn node_runtime(&self) -> Option { match &self.state { AgentServerStoreState::Local { node_runtime, .. } => Some(node_runtime.clone()), _ => None, } } pub fn local( node_runtime: NodeRuntime, fs: Arc, project_environment: Entity, http_client: Arc, cx: &mut Context, ) -> Self { let mut subscriptions = vec![cx.observe_global::(|this, cx| { this.agent_servers_settings_changed(cx); })]; if let Some(registry_store) = AgentRegistryStore::try_global(cx) { subscriptions.push(cx.observe(®istry_store, |this, _, cx| { this.reregister_agents(cx); })); } let mut this = Self { state: AgentServerStoreState::Local { node_runtime, fs, project_environment, http_client, downstream_client: None, settings: None, _subscriptions: subscriptions, }, external_agents: HashMap::default(), }; this.agent_servers_settings_changed(cx); this } pub(crate) fn remote( project_id: u64, upstream_client: Entity, worktree_store: Entity, ) -> Self { Self { state: AgentServerStoreState::Remote { project_id, upstream_client, worktree_store, }, external_agents: HashMap::default(), } } pub fn collab() -> Self { Self { state: AgentServerStoreState::Collab, external_agents: HashMap::default(), } } pub fn shared(&mut self, project_id: u64, client: AnyProtoClient, cx: &mut Context) { match &mut self.state { AgentServerStoreState::Local { downstream_client, .. } => { *downstream_client = Some((project_id, client.clone())); // Send the current list of external agents downstream, but only after a delay, // to avoid having the message arrive before the downstream project's agent server store // sets up its handlers. cx.spawn(async move |this, cx| { cx.background_executor().timer(Duration::from_secs(1)).await; let names = this.update(cx, |this, _| { this.external_agents() .map(|name| name.to_string()) .collect() })?; client .send(proto::ExternalAgentsUpdated { project_id, names }) .log_err(); anyhow::Ok(()) }) .detach(); } AgentServerStoreState::Remote { .. } => { debug_panic!( "external agents over collab not implemented, remote project should not be shared" ); } AgentServerStoreState::Collab => { debug_panic!("external agents over collab not implemented, should not be shared"); } } } pub fn get_external_agent( &mut self, name: &AgentId, ) -> Option<&mut (dyn ExternalAgentServer + 'static)> { self.external_agents .get_mut(name) .map(|entry| entry.server.as_mut()) } pub fn no_browser(&self) -> bool { match &self.state { AgentServerStoreState::Local { downstream_client, .. } => downstream_client .as_ref() .is_some_and(|(_, client)| !client.has_wsl_interop()), _ => false, } } pub fn has_external_agents(&self) -> bool { !self.external_agents.is_empty() } pub fn external_agents(&self) -> impl Iterator { self.external_agents.keys() } async fn handle_get_agent_server_command( this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result { let command = this .update(&mut cx, |this, cx| { let AgentServerStoreState::Local { downstream_client, .. } = &this.state else { debug_panic!("should not receive GetAgentServerCommand in a non-local project"); bail!("unexpected GetAgentServerCommand request in a non-local project"); }; let no_browser = this.no_browser(); let agent = this .external_agents .get_mut(&*envelope.payload.name) .map(|entry| entry.server.as_mut()) .with_context(|| format!("agent `{}` not found", envelope.payload.name))?; let new_version_available_tx = downstream_client .clone() .map(|(project_id, downstream_client)| { let (new_version_available_tx, mut new_version_available_rx) = watch::channel(None); cx.spawn({ let name = envelope.payload.name.clone(); async move |_, _| { if let Some(version) = new_version_available_rx.recv().await.ok().flatten() { downstream_client.send( proto::NewExternalAgentVersionAvailable { project_id, name: name.clone(), version, }, )?; } anyhow::Ok(()) } }) .detach_and_log_err(cx); new_version_available_tx }); let loading_status_tx = downstream_client .clone() .map(|(project_id, downstream_client)| { let (loading_status_tx, mut loading_status_rx) = watch::channel(None); cx.spawn({ let name = envelope.payload.name.clone(); async move |_, _| { while let Ok(status) = loading_status_rx.recv().await { downstream_client.send( proto::ExternalAgentLoadingStatusUpdated { project_id, name: name.clone(), status, }, )?; } anyhow::Ok(()) } }) .detach_and_log_err(cx); loading_status_tx }); let mut extra_env = HashMap::default(); if no_browser { extra_env.insert("NO_BROWSER".to_owned(), "1".to_owned()); } if let Some(new_version_available_tx) = new_version_available_tx { agent.set_new_version_available_tx(new_version_available_tx); } if let Some(loading_status_tx) = loading_status_tx { agent.set_loading_status_tx(loading_status_tx); } anyhow::Ok(agent.get_command(vec![], extra_env, &mut cx.to_async())) })? .await?; Ok(proto::AgentServerCommand { path: command.path.to_string_lossy().into_owned(), args: command.args, env: command .env .map(|env| env.into_iter().collect()) .unwrap_or_default(), root_dir: envelope .payload .root_dir .unwrap_or_else(|| paths::home_dir().to_string_lossy().to_string()), login: None, }) } async fn handle_external_agents_updated( this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result<()> { this.update(&mut cx, |this, cx| { let AgentServerStoreState::Remote { project_id, upstream_client, worktree_store, } = &this.state else { debug_panic!( "handle_external_agents_updated should not be called for a non-remote project" ); bail!("unexpected ExternalAgentsUpdated message") }; let mut previous_entries = std::mem::take(&mut this.external_agents); let mut new_version_available_txs = HashMap::default(); let mut loading_status_txs = HashMap::default(); let mut metadata = HashMap::default(); for (name, mut entry) in previous_entries.drain() { if let Some(tx) = entry.server.take_new_version_available_tx() { new_version_available_txs.insert(name.clone(), tx); } if let Some(tx) = entry.server.take_loading_status_tx() { loading_status_txs.insert(name.clone(), tx); } metadata.insert(name, (entry.icon, entry.display_name, entry.source)); } this.external_agents = envelope .payload .names .into_iter() .map(|name| { let agent_id = AgentId(name.into()); let (icon, display_name, source) = metadata .remove(&agent_id) .or_else(|| { AgentRegistryStore::try_global(cx) .and_then(|store| store.read(cx).agent(&agent_id)) .map(|s| { ( s.icon_path().cloned(), Some(s.name().clone()), ExternalAgentSource::Registry, ) }) }) .unwrap_or((None, None, ExternalAgentSource::default())); let agent = RemoteExternalAgentServer { project_id: *project_id, upstream_client: upstream_client.clone(), worktree_store: worktree_store.clone(), name: agent_id.clone(), new_version_available_tx: new_version_available_txs.remove(&agent_id), loading_status_tx: loading_status_txs.remove(&agent_id), }; ( agent_id, ExternalAgentEntry::new( Box::new(agent) as Box, source, icon, display_name, ), ) }) .collect(); cx.emit(AgentServersUpdated); Ok(()) }) } async fn handle_loading_status_updated( this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result<()> { this.update(&mut cx, |this, _| { if let Some(entry) = this.external_agents.get_mut(&*envelope.payload.name) && let Some(mut tx) = entry.server.take_loading_status_tx() { tx.send(envelope.payload.status).ok(); entry.server.set_loading_status_tx(tx); } }); Ok(()) } async fn handle_new_version_available( this: Entity, envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result<()> { this.update(&mut cx, |this, _| { if let Some(entry) = this.external_agents.get_mut(&*envelope.payload.name) && let Some(mut tx) = entry.server.take_new_version_available_tx() { tx.send(Some(envelope.payload.version)).ok(); entry.server.set_new_version_available_tx(tx); } }); Ok(()) } } struct RemoteExternalAgentServer { project_id: u64, upstream_client: Entity, worktree_store: Entity, name: AgentId, new_version_available_tx: Option>>, loading_status_tx: Option>>, } impl ExternalAgentServer for RemoteExternalAgentServer { fn take_new_version_available_tx(&mut self) -> Option>> { self.new_version_available_tx.take() } fn set_new_version_available_tx(&mut self, tx: watch::Sender>) { self.new_version_available_tx = Some(tx); } fn take_loading_status_tx(&mut self) -> Option>> { self.loading_status_tx.take() } fn set_loading_status_tx(&mut self, tx: watch::Sender>) { self.loading_status_tx = Some(tx); } fn get_command( &mut self, extra_args: Vec, extra_env: HashMap, cx: &mut AsyncApp, ) -> Task> { let project_id = self.project_id; let name = self.name.to_string(); let upstream_client = self.upstream_client.downgrade(); let worktree_store = self.worktree_store.clone(); cx.spawn(async move |cx| { let root_dir = worktree_store.read_with(cx, |worktree_store, cx| { crate::Project::default_visible_worktree_paths(worktree_store, cx) .into_iter() .next() .map(|path| path.display().to_string()) }); let mut response = upstream_client .update(cx, |upstream_client, _| { upstream_client .proto_client() .request(proto::GetAgentServerCommand { project_id, name, root_dir, }) })? .await?; response.args.extend(extra_args); response.env.extend(extra_env); Ok(AgentServerCommand { path: response.path.into(), args: response.args, env: Some(response.env.into_iter().collect()), }) }) } fn as_any(&self) -> &dyn Any { self } fn as_any_mut(&mut self) -> &mut dyn Any { self } } #[derive(Debug, PartialEq, Eq)] enum RegistryArchiveKind { Archive(AssetKind), /// The archive URL points directly at an executable, per the ACP registry /// schema: "URL to download archive (.zip, .tar.gz, .tgz, .tar.bz2, .tbz2, /// or raw binary)". RawBinary { file_name: String, }, } fn registry_archive_kind_for_url(archive_url: &str) -> Result { const UNSUPPORTED_SUFFIXES: &[&str] = &[ // Installer formats explicitly rejected by the registry schema. ".dmg", ".pkg", ".deb", ".rpm", ".msi", ".appimage", // Archive formats we cannot extract; treating them as raw binaries // would produce a broken install. ".tar.xz", ".txz", ".tar", ".gz", ".bz2", ".xz", ".7z", ]; let archive_path = Url::parse(archive_url) .ok() .map(|url| url.path().to_string()) .unwrap_or_else(|| archive_url.to_string()); let lowercase_path = archive_path.to_lowercase(); if lowercase_path.ends_with(".zip") { Ok(RegistryArchiveKind::Archive(AssetKind::Zip)) } else if lowercase_path.ends_with(".tar.gz") || lowercase_path.ends_with(".tgz") { Ok(RegistryArchiveKind::Archive(AssetKind::TarGz)) } else if lowercase_path.ends_with(".tar.bz2") || lowercase_path.ends_with(".tbz2") { Ok(RegistryArchiveKind::Archive(AssetKind::TarBz2)) } else if let Some(suffix) = UNSUPPORTED_SUFFIXES .iter() .find(|suffix| lowercase_path.ends_with(*suffix)) { bail!("unsupported archive type {suffix} in URL: {archive_url}"); } else { let file_name = raw_binary_file_name(&archive_path) .with_context(|| format!("determining binary file name from URL: {archive_url}"))?; Ok(RegistryArchiveKind::RawBinary { file_name }) } } fn raw_binary_file_name(archive_path: &str) -> Result { let last_segment = archive_path .rsplit('/') .next() .filter(|segment| !segment.is_empty()) .context("URL has no file name")?; let file_name = percent_decode_str(last_segment) .decode_utf8() .context("file name is not valid UTF-8")? .into_owned(); anyhow::ensure!( !file_name.is_empty() && file_name != "." && file_name != ".." && !file_name.contains(['/', '\\']) && !file_name.contains('\0'), "invalid binary file name: {file_name}" ); Ok(file_name) } struct GithubReleaseArchive { repo_name_with_owner: String, tag: String, asset_name: String, } fn github_release_archive_from_url(archive_url: &str) -> Option { fn decode_path_segment(segment: &str) -> Option { percent_decode_str(segment) .decode_utf8() .ok() .map(|segment| segment.into_owned()) } let url = Url::parse(archive_url).ok()?; if url.scheme() != "https" || url.host_str()? != "github.com" { return None; } let segments = url.path_segments()?.collect::>(); if segments.len() < 6 || segments[2] != "releases" || segments[3] != "download" { return None; } Some(GithubReleaseArchive { repo_name_with_owner: format!("{}/{}", segments[0], segments[1]), tag: decode_path_segment(segments[4])?, asset_name: segments[5..] .iter() .map(|segment| decode_path_segment(segment)) .collect::>>()? .join("/"), }) } fn sanitize_path_component(input: &str) -> String { let sanitized = input .chars() .map(|character| match character { 'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '_' | '-' => character, _ => '-', }) .collect::(); if sanitized.is_empty() { "unknown".to_string() } else { sanitized } } fn versioned_archive_cache_dir( base_dir: &Path, version: Option<&str>, archive_url: &str, sha256: Option<&str>, ) -> PathBuf { let version = version.unwrap_or_default(); let sanitized_version = sanitize_path_component(version); let mut version_hasher = Sha256::new(); version_hasher.update(version.as_bytes()); let version_hash = format!("{:x}", version_hasher.finalize()); let mut archive_hasher = Sha256::new(); archive_hasher.update(archive_url.as_bytes()); if let Some(sha256) = sha256 { archive_hasher.update(b"\0sha256:"); archive_hasher.update(sha256.to_ascii_lowercase().as_bytes()); } let archive_hash = format!("{:x}", archive_hasher.finalize()); base_dir.join(format!( "v_{sanitized_version}_{}_{}", &version_hash[..16], &archive_hash[..16], )) } // The `v_` prefix here must stay in sync with `versioned_archive_cache_dir`, // so we only ever remove directories that we created ourselves. const VERSIONED_ARCHIVE_CACHE_DIR_PREFIX: &str = "v_"; async fn remove_stale_versioned_archive_cache_dirs( fs: Arc, base_dir: &Path, current_version_dir: &Path, ) -> Result<()> { let Some(current_dir_name) = current_version_dir.file_name() else { return Ok(()); }; let current_mtime = fs .metadata(current_version_dir) .await .with_context(|| format!("reading metadata for {current_version_dir:?}"))? .with_context(|| format!("missing metadata for {current_version_dir:?}"))? .mtime; let mut entries = fs .read_dir(base_dir) .await .with_context(|| format!("reading archive cache directory {base_dir:?}"))?; while let Some(entry) = entries.next().await { let entry = entry.with_context(|| format!("reading entry in {base_dir:?}"))?; let Some(entry_name) = entry.file_name() else { continue; }; if entry_name == current_dir_name || !entry_name .to_string_lossy() .starts_with(VERSIONED_ARCHIVE_CACHE_DIR_PREFIX) { continue; } let Some(entry_metadata) = fs.metadata(&entry).await.log_err().flatten() else { continue; }; if !entry_metadata.is_dir { continue; } // Only remove directories that predate the current version's directory. // This avoids racing with a concurrent extraction of a different version // that finished after we cached the current version's mtime. if !current_mtime.bad_is_greater_than(entry_metadata.mtime) { continue; } fs.remove_dir( &entry, RemoveOptions { recursive: true, ignore_if_not_exists: true, }, ) .await .with_context(|| format!("removing stale archive cache directory {entry:?}"))?; } Ok(()) } struct LocalRegistryArchiveAgent { fs: Arc, http_client: Arc, node_runtime: NodeRuntime, project_environment: Entity, installation_dir: PathBuf, version: SharedString, targets: HashMap, env: HashMap, new_version_available_tx: Option>>, loading_status_tx: Option>>, } impl ExternalAgentServer for LocalRegistryArchiveAgent { fn version(&self) -> Option<&SharedString> { Some(&self.version) } fn take_new_version_available_tx(&mut self) -> Option>> { self.new_version_available_tx.take() } fn set_new_version_available_tx(&mut self, tx: watch::Sender>) { self.new_version_available_tx = Some(tx); } fn take_loading_status_tx(&mut self) -> Option>> { self.loading_status_tx.take() } fn set_loading_status_tx(&mut self, tx: watch::Sender>) { self.loading_status_tx = Some(tx); } fn get_command( &mut self, extra_args: Vec, extra_env: HashMap, cx: &mut AsyncApp, ) -> Task> { let fs = self.fs.clone(); let http_client = self.http_client.clone(); let node_runtime = self.node_runtime.clone(); let project_environment = self.project_environment.downgrade(); let installation_dir = self.installation_dir.clone(); let targets = self.targets.clone(); let settings_env = self.env.clone(); let version = self.version.clone(); let loading_status_tx = self.loading_status_tx.take(); cx.spawn(async move |cx| { let mut env = project_environment .update(cx, |project_environment, cx| { project_environment.default_environment(cx) })? .await .unwrap_or_default(); let dir = installation_dir; fs.create_dir(&dir).await?; let os = if cfg!(target_os = "macos") { "darwin" } else if cfg!(target_os = "linux") { "linux" } else if cfg!(target_os = "windows") { "windows" } else { anyhow::bail!("unsupported OS"); }; let arch = if cfg!(target_arch = "aarch64") { "aarch64" } else if cfg!(target_arch = "x86_64") { "x86_64" } else { anyhow::bail!("unsupported architecture"); }; let platform_key = format!("{}-{}", os, arch); let target_config = targets.get(&platform_key).with_context(|| { format!( "no target specified for platform '{}'. Available platforms: {}", platform_key, targets .keys() .map(|k| k.as_str()) .collect::>() .join(", ") ) })?; env.extend(target_config.env.clone()); env.extend(extra_env); env.extend(settings_env); let archive_url = &target_config.archive; let version_dir = versioned_archive_cache_dir( &dir, Some(version.as_ref()), archive_url, target_config.sha256.as_deref(), ); if !fs.is_dir(&version_dir).await { let mut loading_status_tx = loading_status_tx; if let Some(tx) = loading_status_tx.as_mut() { tx.send(Some(format!("Installing {}…", version.as_ref()))) .ok(); } let sha256 = if let Some(provided_sha) = &target_config.sha256 { Some(provided_sha.clone()) } else if let Some(github_archive) = github_release_archive_from_url(archive_url) { if let Ok(release) = ::http_client::github::get_release_by_tag_name( &github_archive.repo_name_with_owner, &github_archive.tag, http_client.clone(), ) .await { if let Some(asset) = release .assets .iter() .find(|a| a.name == github_archive.asset_name) { asset.digest.as_ref().and_then(|d| { d.strip_prefix("sha256:") .map(|s| s.to_string()) .or_else(|| Some(d.clone())) }) } else { None } } else { None } } else { None }; match registry_archive_kind_for_url(archive_url)? { RegistryArchiveKind::Archive(asset_kind) => { ::http_client::github_download::download_server_binary( &*http_client, archive_url, sha256.as_deref(), &version_dir, asset_kind, ) .await?; } RegistryArchiveKind::RawBinary { file_name } => { ::http_client::github_download::download_server_raw_binary( &*http_client, archive_url, sha256.as_deref(), &version_dir, &file_name, ) .await?; } } } let cmd = &target_config.cmd; let cmd_path = if cmd == "node" { node_runtime.binary_path().await? } else { if cmd.contains("..") { anyhow::bail!("command path cannot contain '..': {}", cmd); } if cmd.starts_with("./") || cmd.starts_with(".\\") { let cmd_path = version_dir.join(&cmd[2..]); anyhow::ensure!( fs.is_file(&cmd_path).await, "Missing command {} after extraction", cmd_path.to_string_lossy() ); cmd_path } else { anyhow::bail!("command must be relative (start with './'): {}", cmd); } }; cx.background_spawn({ let fs = fs.clone(); let dir = dir.clone(); let version_dir = version_dir.clone(); async move { remove_stale_versioned_archive_cache_dirs(fs, &dir, &version_dir) .await .log_err(); } }) .detach(); let mut args = target_config.args.clone(); args.extend(extra_args); let command = AgentServerCommand { path: cmd_path, args, env: Some(env), }; Ok(command) }) } fn as_any(&self) -> &dyn Any { self } fn as_any_mut(&mut self) -> &mut dyn Any { self } } struct LocalRegistryNpxAgent { fs: Arc, node_runtime: NodeRuntime, project_environment: Entity, registry_id: Arc, version: SharedString, package: SharedString, args: Vec, distribution_env: HashMap, settings_env: HashMap, new_version_available_tx: Option>>, } impl ExternalAgentServer for LocalRegistryNpxAgent { fn version(&self) -> Option<&SharedString> { Some(&self.version) } fn take_new_version_available_tx(&mut self) -> Option>> { self.new_version_available_tx.take() } fn set_new_version_available_tx(&mut self, tx: watch::Sender>) { self.new_version_available_tx = Some(tx); } fn get_command( &mut self, extra_args: Vec, extra_env: HashMap, cx: &mut AsyncApp, ) -> Task> { let fs = self.fs.clone(); let node_runtime = self.node_runtime.clone(); let project_environment = self.project_environment.downgrade(); let registry_id = self.registry_id.clone(); let package = self.package.clone(); let args = self.args.clone(); let distribution_env = self.distribution_env.clone(); let settings_env = self.settings_env.clone(); cx.spawn(async move |cx| { let mut env = project_environment .update(cx, |project_environment, cx| { project_environment.default_environment(cx) })? .await .unwrap_or_default(); let install_dir = paths::external_agents_dir() .join("registry") .join("npx") .join(sanitize_path_component(®istry_id)); fs.create_dir(&install_dir).await?; let (package_name, package_spec) = bounded_npm_package_spec(&package); node_runtime .run_npm_subcommand( Some(&install_dir), "install", &[package_spec.as_str(), "--save-exact"], ) .await?; let executable = node_runtime::read_package_executable( install_dir.join("node_modules"), package_name, ) .await?; let node_binary = node_runtime.binary_path().await?; env.extend(node_runtime::npm_command_env(&node_binary)); env.extend(distribution_env); env.extend(extra_env); env.extend(settings_env); let mut command_args = vec![executable.to_string_lossy().into_owned()]; command_args.extend(args); command_args.extend(extra_args); let command = AgentServerCommand { path: node_binary, args: command_args, env: Some(env), }; Ok(command) }) } fn as_any(&self) -> &dyn Any { self } fn as_any_mut(&mut self) -> &mut dyn Any { self } } /// People are using min-release-age more frequently. Which means a fresh registry will likely have /// new package versions than the user can install. /// We set the version to now be a ceiling and not an exact pin instead. This allows npm to resolve /// the latest version it can find that satisfies the constraint. npm seems to check regularly enough /// that new versions are available. This does have a few downsides: /// - The user might have an older cached version of the package that satisfies the constraint, until /// npm checks for updates again. /// - The registry args/env may not be valid for the resolved version. /// /// This is a best-effort attempt to install a version that works without overriding the user's /// security settings, as the args don't change often. The registry will need to support this better /// at some point, but until then, this is a best-effort workaround that hopefully solves the issue /// for most users. /// /// We use npm's hyphen-range syntax (`0.0.0 - `, equivalent to `<=`) instead of /// the more compact `<=` form because on Windows, `npm` is `npm.cmd` (a batch file run by /// cmd.exe), and the quotes our shell builder emits are PowerShell string-literal syntax that PS /// strips during parsing. PS only re-adds CRT-style transport quotes around native command args /// containing whitespace, so `package@<=0.25.3` reaches cmd.exe bare and the unquoted `<` is /// interpreted as input redirection. See zed-industries/zed#55921. fn bounded_npm_package_spec(package_spec: &str) -> (&str, String) { let Some((package_name, version)) = package_spec.rsplit_once('@') else { return (package_spec, package_spec.to_string()); }; if package_name.is_empty() { return (package_spec, package_spec.to_string()); } if Version::parse(version).is_err() { return (package_name, package_spec.to_string()); } (package_name, format!("{package_name}@0.0.0 - {version}")) } struct LocalCustomAgent { project_environment: Entity, command: AgentServerCommand, } impl ExternalAgentServer for LocalCustomAgent { fn get_command( &mut self, extra_args: Vec, extra_env: HashMap, cx: &mut AsyncApp, ) -> Task> { let mut command = self.command.clone(); let project_environment = self.project_environment.downgrade(); cx.spawn(async move |cx| { let mut env = project_environment .update(cx, |project_environment, cx| { project_environment.default_environment(cx) })? .await .unwrap_or_default(); env.extend(command.env.unwrap_or_default()); env.extend(extra_env); command.env = Some(env); command.args.extend(extra_args); Ok(command) }) } fn as_any(&self) -> &dyn Any { self } fn as_any_mut(&mut self) -> &mut dyn Any { self } } #[derive(Default, Clone, JsonSchema, Debug, PartialEq, RegisterSetting)] pub struct AllAgentServersSettings(pub HashMap); impl std::ops::Deref for AllAgentServersSettings { type Target = HashMap; fn deref(&self) -> &Self::Target { &self.0 } } impl std::ops::DerefMut for AllAgentServersSettings { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } impl AllAgentServersSettings { pub fn has_registry_agents(&self) -> bool { self.values() .any(|s| matches!(s, CustomAgentServerSettings::Registry { .. })) } } #[derive(Clone, JsonSchema, Debug, PartialEq)] pub enum CustomAgentServerSettings { Custom { command: AgentServerCommand, /// The default mode to use for this agent. /// /// Note: Not only all agents support modes. /// /// Default: None default_mode: Option, /// Default values for session config options. /// /// This is a map from config option ID to the default value for that option. /// /// Default: {} default_config_options: HashMap, /// Favorited values for session config options. /// /// This is a map from config option ID to a list of favorited value IDs. /// /// Default: {} favorite_config_option_values: HashMap>, }, Registry { /// Additional environment variables to pass to the agent. /// /// Default: {} env: HashMap, /// The default mode to use for this agent. /// /// Note: Not only all agents support modes. /// /// Default: None default_mode: Option, /// Default values for session config options. /// /// This is a map from config option ID to the default value for that option. /// /// Default: {} default_config_options: HashMap, /// Favorited values for session config options. /// /// This is a map from config option ID to a list of favorited value IDs. /// /// Default: {} favorite_config_option_values: HashMap>, }, } impl CustomAgentServerSettings { pub fn command(&self) -> Option<&AgentServerCommand> { match self { CustomAgentServerSettings::Custom { command, .. } => Some(command), CustomAgentServerSettings::Registry { .. } => None, } } pub fn default_mode(&self) -> Option<&str> { match self { CustomAgentServerSettings::Custom { default_mode, .. } | CustomAgentServerSettings::Registry { default_mode, .. } => default_mode.as_deref(), } } pub fn default_config_option(&self, config_id: &str) -> Option<&AgentConfigOptionValue> { match self { CustomAgentServerSettings::Custom { default_config_options, .. } | CustomAgentServerSettings::Registry { default_config_options, .. } => default_config_options.get(config_id), } } pub fn favorite_config_option_values(&self, config_id: &str) -> Option<&[String]> { match self { CustomAgentServerSettings::Custom { favorite_config_option_values, .. } | CustomAgentServerSettings::Registry { favorite_config_option_values, .. } => favorite_config_option_values .get(config_id) .map(|v| v.as_slice()), } } } impl From for CustomAgentServerSettings { fn from(value: settings::CustomAgentServerSettings) -> Self { match value { settings::CustomAgentServerSettings::Custom { path, args, env, default_mode, default_config_options, favorite_config_option_values, } => CustomAgentServerSettings::Custom { command: AgentServerCommand { path: PathBuf::from(shellexpand::tilde(&path.to_string_lossy()).as_ref()), args, env: Some(env), }, default_mode, default_config_options, favorite_config_option_values, }, settings::CustomAgentServerSettings::Registry { env, default_mode, default_config_options, favorite_config_option_values, } => CustomAgentServerSettings::Registry { env, default_mode, default_config_options, favorite_config_option_values, }, } } } impl settings::Settings for AllAgentServersSettings { fn from_settings(content: &settings::SettingsContent) -> Self { let agent_settings = content.agent_servers.clone().unwrap(); Self( agent_settings .0 .into_iter() .map(|(k, v)| { ( EXTENSION_TO_REGISTRY_IDS .get(&k.as_str()) .map(|v| v.to_string()) .unwrap_or(k), v.into(), ) }) .collect(), ) } } #[cfg(test)] mod tests { use super::*; use crate::agent_registry_store::{ AgentRegistryStore, RegistryAgent, RegistryAgentMetadata, RegistryNpxAgent, }; use crate::worktree_store::{WorktreeIdCounter, WorktreeStore}; use gpui::TestAppContext; #[cfg(feature = "test-support")] use http_client::{AsyncBody, FakeHttpClient, Response}; use node_runtime::NodeRuntime; use settings::Settings as _; #[cfg(feature = "test-support")] const TEST_ARCHIVE_URL: &str = "https://example.test/agent"; #[cfg(feature = "test-support")] fn static_http_client(body: Vec) -> Arc { FakeHttpClient::create(move |_| { let body = body.clone(); async move { Ok(Response::builder() .status(200) .body(AsyncBody::from(body))?) } }) } #[cfg(feature = "test-support")] fn make_registry_archive_agent( cx: &mut TestAppContext, installation_dir: PathBuf, http_client: Arc, sha256: Option, ) -> LocalRegistryArchiveAgent { let fs: Arc = Arc::new(fs::RealFs::new(None, cx.executor())); let target = RegistryTargetConfig { archive: TEST_ARCHIVE_URL.to_string(), cmd: "./agent".to_string(), args: Vec::new(), sha256, env: HashMap::default(), }; let targets = [ "darwin-aarch64", "darwin-x86_64", "linux-aarch64", "linux-x86_64", "windows-aarch64", "windows-x86_64", ] .into_iter() .map(|platform| (platform.to_string(), target.clone())) .collect(); cx.update(|cx| { let worktree_store = cx.new(|cx| WorktreeStore::local(false, fs.clone(), WorktreeIdCounter::get(cx))); let project_environment = cx.new(|cx| { crate::ProjectEnvironment::new(None, worktree_store.downgrade(), None, false, cx) }); LocalRegistryArchiveAgent { fs, http_client, node_runtime: NodeRuntime::unavailable(), project_environment, installation_dir, version: "1.0.0".into(), targets, env: HashMap::default(), new_version_available_tx: None, loading_status_tx: None, } }) } fn make_npx_agent(id: &str, version: &str) -> RegistryAgent { let id = SharedString::from(id.to_string()); RegistryAgent::Npx(RegistryNpxAgent { metadata: RegistryAgentMetadata { id: AgentId::new(id.clone()), name: id.clone(), description: SharedString::from(""), version: SharedString::from(version.to_string()), repository: None, website: None, icon_path: None, }, package: id, args: Vec::new(), env: HashMap::default(), }) } fn init_test_settings(cx: &mut TestAppContext) { cx.update(|cx| { let settings_store = SettingsStore::test(cx); cx.set_global(settings_store); }); } fn init_registry( cx: &mut TestAppContext, agents: Vec, ) -> gpui::Entity { cx.update(|cx| AgentRegistryStore::init_test_global(cx, agents)) } fn set_registry_settings(cx: &mut TestAppContext, agent_names: &[&str]) { cx.update(|cx| { AllAgentServersSettings::override_global( AllAgentServersSettings( agent_names .iter() .map(|name| { ( name.to_string(), settings::CustomAgentServerSettings::Registry { env: HashMap::default(), default_mode: None, default_config_options: HashMap::default(), favorite_config_option_values: HashMap::default(), } .into(), ) }) .collect(), ), cx, ); }); } fn create_agent_server_store(cx: &mut TestAppContext) -> gpui::Entity { cx.update(|cx| { let fs: Arc = fs::FakeFs::new(cx.background_executor().clone()); let worktree_store = cx.new(|cx| WorktreeStore::local(false, fs.clone(), WorktreeIdCounter::get(cx))); let project_environment = cx.new(|cx| { crate::ProjectEnvironment::new(None, worktree_store.downgrade(), None, false, cx) }); let http_client = http_client::FakeHttpClient::with_404_response(); cx.new(|cx| { AgentServerStore::local( NodeRuntime::unavailable(), fs, project_environment, http_client, cx, ) }) }) } #[test] fn builds_bounded_npm_package_specs() { assert_eq!( bounded_npm_package_spec("agent-package@1.2.3"), ("agent-package", "agent-package@0.0.0 - 1.2.3".to_string()) ); assert_eq!( bounded_npm_package_spec("@scope/agent-package@1.2.3-beta.1"), ( "@scope/agent-package", "@scope/agent-package@0.0.0 - 1.2.3-beta.1".to_string() ) ); assert_eq!( bounded_npm_package_spec("@scope/agent-package"), ("@scope/agent-package", "@scope/agent-package".to_string()) ); assert_eq!( bounded_npm_package_spec("agent-package@latest"), ("agent-package", "agent-package@latest".to_string()) ); } #[test] fn detects_supported_archive_suffixes() { assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.zip"), Ok(RegistryArchiveKind::Archive(AssetKind::Zip)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.zip?download=1"), Ok(RegistryArchiveKind::Archive(AssetKind::Zip)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tar.gz"), Ok(RegistryArchiveKind::Archive(AssetKind::TarGz)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tar.gz?download=1#latest"), Ok(RegistryArchiveKind::Archive(AssetKind::TarGz)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tgz"), Ok(RegistryArchiveKind::Archive(AssetKind::TarGz)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tgz#download"), Ok(RegistryArchiveKind::Archive(AssetKind::TarGz)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tar.bz2"), Ok(RegistryArchiveKind::Archive(AssetKind::TarBz2)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tar.bz2?download=1"), Ok(RegistryArchiveKind::Archive(AssetKind::TarBz2)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tbz2"), Ok(RegistryArchiveKind::Archive(AssetKind::TarBz2)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.tbz2#download"), Ok(RegistryArchiveKind::Archive(AssetKind::TarBz2)) )); assert!(matches!( registry_archive_kind_for_url("https://example.com/agent.ZIP"), Ok(RegistryArchiveKind::Archive(AssetKind::Zip)) )); } #[test] fn detects_raw_binary_archive_urls() { assert_eq!( registry_archive_kind_for_url("https://x.ai/cli/grok-0.2.20-macos-aarch64").unwrap(), RegistryArchiveKind::RawBinary { file_name: "grok-0.2.20-macos-aarch64".to_string() }, ); assert_eq!( registry_archive_kind_for_url("https://x.ai/cli/grok-0.2.20-windows-x86_64.exe") .unwrap(), RegistryArchiveKind::RawBinary { file_name: "grok-0.2.20-windows-x86_64.exe".to_string() }, ); assert_eq!( registry_archive_kind_for_url("https://example.com/agent-binary?download=1#latest") .unwrap(), RegistryArchiveKind::RawBinary { file_name: "agent-binary".to_string() }, ); assert_eq!( registry_archive_kind_for_url("https://example.com/agent%20binary").unwrap(), RegistryArchiveKind::RawBinary { file_name: "agent binary".to_string() }, ); // No file name to install the binary as. assert!(registry_archive_kind_for_url("https://example.com/").is_err()); // Percent-decoding must not allow path traversal in the file name. assert!(registry_archive_kind_for_url("https://example.com/a%2F..%2Fevil").is_err()); assert!(registry_archive_kind_for_url("https://example.com/%2E%2E").is_err()); } #[test] fn parses_github_release_archive_urls() { let github_archive = github_release_archive_from_url( "https://github.com/owner/repo/releases/download/release%2F2.3.5/agent.tar.bz2?download=1", ) .unwrap(); assert_eq!(github_archive.repo_name_with_owner, "owner/repo"); assert_eq!(github_archive.tag, "release/2.3.5"); assert_eq!(github_archive.asset_name, "agent.tar.bz2"); } #[test] fn rejects_unsupported_archive_suffixes() { let error = registry_archive_kind_for_url("https://example.com/agent.tar.xz") .err() .map(|error| error.to_string()); assert_eq!( error, Some( "unsupported archive type .tar.xz in URL: https://example.com/agent.tar.xz" .to_string() ), ); for installer_url in [ "https://example.com/agent.dmg", "https://example.com/agent.pkg", "https://example.com/agent.deb", "https://example.com/agent.rpm", "https://example.com/agent.msi", "https://example.com/agent.AppImage", ] { assert!( registry_archive_kind_for_url(installer_url).is_err(), "expected {installer_url} to be rejected" ); } } #[test] fn versioned_archive_cache_dir_includes_artifact_identity() { let slash_version_dir = versioned_archive_cache_dir( Path::new("/tmp/agents"), Some("release/2.3.5"), "https://example.com/agent.zip", None, ); let colon_version_dir = versioned_archive_cache_dir( Path::new("/tmp/agents"), Some("release:2.3.5"), "https://example.com/agent.zip", None, ); let file_name = slash_version_dir .file_name() .and_then(|name| name.to_str()) .expect("cache directory should have a file name"); assert!(file_name.starts_with("v_release-2.3.5_")); assert_ne!(slash_version_dir, colon_version_dir); let lowercase_checksum_dir = versioned_archive_cache_dir( Path::new("/tmp/agents"), Some("release/2.3.5"), "https://example.com/agent.zip", Some("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), ); let uppercase_checksum_dir = versioned_archive_cache_dir( Path::new("/tmp/agents"), Some("release/2.3.5"), "https://example.com/agent.zip", Some("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"), ); let changed_checksum_dir = versioned_archive_cache_dir( Path::new("/tmp/agents"), Some("release/2.3.5"), "https://example.com/agent.zip", Some("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"), ); assert_ne!(slash_version_dir, lowercase_checksum_dir); assert_eq!(lowercase_checksum_dir, uppercase_checksum_dir); assert_ne!(lowercase_checksum_dir, changed_checksum_dir); } #[cfg(feature = "test-support")] #[gpui::test] async fn registry_raw_binary_checksum_invalidates_unverified_cache_and_blocks_mismatch( cx: &mut TestAppContext, ) { init_test_settings(cx); cx.executor().allow_parking(); let temp_dir = tempfile::tempdir().unwrap(); let installation_dir = temp_dir.path().join("agent"); let old_version_dir = versioned_archive_cache_dir(&installation_dir, Some("1.0.0"), TEST_ARCHIVE_URL, None); std::fs::create_dir_all(&old_version_dir).unwrap(); std::fs::write(old_version_dir.join("agent"), b"unverified agent").unwrap(); let expected_sha256 = "0000000000000000000000000000000000000000000000000000000000000000"; let http_client = static_http_client(b"unexpected agent".to_vec()); let mut agent = make_registry_archive_agent( cx, installation_dir.clone(), http_client, Some(expected_sha256.to_string()), ); let get_command = cx.update(|cx| agent.get_command(Vec::new(), HashMap::default(), &mut cx.to_async())); let error = get_command.await.unwrap_err(); assert!( error.to_string().contains("SHA-256 mismatch"), "unexpected error: {error:#}" ); assert!(old_version_dir.exists()); assert!( !versioned_archive_cache_dir( &installation_dir, Some("1.0.0"), TEST_ARCHIVE_URL, Some(expected_sha256), ) .exists() ); } #[cfg(feature = "test-support")] #[gpui::test] async fn registry_raw_binary_with_checksum_installs(cx: &mut TestAppContext) { init_test_settings(cx); cx.executor().allow_parking(); let temp_dir = tempfile::tempdir().unwrap(); let installation_dir = temp_dir.path().join("agent"); let contents = b"verified agent"; let expected_sha256 = format!("{:X}", Sha256::digest(contents)); let http_client = static_http_client(contents.to_vec()); let mut agent = make_registry_archive_agent( cx, installation_dir.clone(), http_client, Some(expected_sha256.clone()), ); let get_command = cx.update(|cx| agent.get_command(Vec::new(), HashMap::default(), &mut cx.to_async())); let command = get_command.await.unwrap(); cx.run_until_parked(); assert_eq!( command.path, versioned_archive_cache_dir( &installation_dir, Some("1.0.0"), TEST_ARCHIVE_URL, Some(&expected_sha256), ) .join("agent") ); assert_eq!(std::fs::read(command.path).unwrap(), contents); } #[cfg(feature = "test-support")] #[gpui::test] async fn registry_raw_binary_without_checksum_installs(cx: &mut TestAppContext) { init_test_settings(cx); cx.executor().allow_parking(); let temp_dir = tempfile::tempdir().unwrap(); let installation_dir = temp_dir.path().join("agent"); let contents = b"unchecked agent"; let http_client = static_http_client(contents.to_vec()); let mut agent = make_registry_archive_agent(cx, installation_dir.clone(), http_client, None); let get_command = cx.update(|cx| agent.get_command(Vec::new(), HashMap::default(), &mut cx.to_async())); let command = get_command.await.unwrap(); cx.run_until_parked(); assert_eq!( command.path, versioned_archive_cache_dir(&installation_dir, Some("1.0.0"), TEST_ARCHIVE_URL, None,) .join("agent") ); assert_eq!(std::fs::read(command.path).unwrap(), contents); } #[gpui::test] async fn test_remove_stale_versioned_archive_cache_dirs(cx: &mut TestAppContext) { let fs = fs::FakeFs::new(cx.executor()); let base_dir = Path::new("/cache"); // FakeFs increments mtime on every create, so creation order is // ascending mtime: v_old_1 < v_old_2 < other < v_not_a_dir < v_current < v_newer. fs.insert_tree( base_dir, serde_json::json!({ "v_old_1": {}, "v_old_2": {}, "other": {}, }), ) .await; fs.insert_file(base_dir.join("v_not_a_dir"), b"keep me".to_vec()) .await; let current_version_dir = base_dir.join("v_current"); fs.create_dir(¤t_version_dir).await.unwrap(); // Sibling that "finished extracting" after the current dir was cached. fs.create_dir(&base_dir.join("v_newer")).await.unwrap(); remove_stale_versioned_archive_cache_dirs( fs.clone() as Arc, base_dir, ¤t_version_dir, ) .await .unwrap(); let mut remaining = fs .read_dir(base_dir) .await .unwrap() .filter_map(|entry| async move { entry.ok() }) .map(|path| { path.file_name() .expect("entry has a name") .to_string_lossy() .into_owned() }) .collect::>() .await; remaining.sort(); assert_eq!( remaining, vec![ "other".to_string(), "v_current".to_string(), "v_newer".to_string(), "v_not_a_dir".to_string(), ] ); } #[gpui::test] fn test_version_change_sends_notification(cx: &mut TestAppContext) { init_test_settings(cx); let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]); set_registry_settings(cx, &["test-agent"]); let store = create_agent_server_store(cx); // Verify the agent was registered with version 1.0.0. store.read_with(cx, |store, _| { let entry = store .external_agents .get(&AgentId::new("test-agent")) .expect("agent should be registered"); assert_eq!( entry.server.version().map(|v| v.to_string()), Some("1.0.0".to_string()) ); }); // Set up a watch channel and store the tx on the agent. let (tx, mut rx) = watch::channel::>(None); store.update(cx, |store, _| { let entry = store .external_agents .get_mut(&AgentId::new("test-agent")) .expect("agent should be registered"); entry.server.set_new_version_available_tx(tx); }); // Update the registry to version 2.0.0. registry.update(cx, |store, cx| { store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx); }); cx.run_until_parked(); // The watch channel should have received the new version. assert_eq!(rx.borrow().as_deref(), Some("2.0.0")); } #[gpui::test] fn test_same_version_preserves_tx(cx: &mut TestAppContext) { init_test_settings(cx); let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]); set_registry_settings(cx, &["test-agent"]); let store = create_agent_server_store(cx); let (tx, mut rx) = watch::channel::>(None); store.update(cx, |store, _| { let entry = store .external_agents .get_mut(&AgentId::new("test-agent")) .expect("agent should be registered"); entry.server.set_new_version_available_tx(tx); }); // "Refresh" the registry with the same version. registry.update(cx, |store, cx| { store.set_agents(vec![make_npx_agent("test-agent", "1.0.0")], cx); }); cx.run_until_parked(); // No notification should have been sent. assert_eq!(rx.borrow().as_deref(), None); // The tx should have been transferred to the rebuilt agent entry. store.update(cx, |store, _| { let entry = store .external_agents .get_mut(&AgentId::new("test-agent")) .expect("agent should be registered"); assert!( entry.server.take_new_version_available_tx().is_some(), "tx should have been transferred to the rebuilt agent" ); }); } #[gpui::test] fn test_no_tx_stored_does_not_panic_on_version_change(cx: &mut TestAppContext) { init_test_settings(cx); let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]); set_registry_settings(cx, &["test-agent"]); let _store = create_agent_server_store(cx); // Update the registry without having stored any tx — should not panic. registry.update(cx, |store, cx| { store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx); }); cx.run_until_parked(); } #[gpui::test] fn test_multiple_agents_independent_notifications(cx: &mut TestAppContext) { init_test_settings(cx); let registry = init_registry( cx, vec![ make_npx_agent("agent-a", "1.0.0"), make_npx_agent("agent-b", "3.0.0"), ], ); set_registry_settings(cx, &["agent-a", "agent-b"]); let store = create_agent_server_store(cx); let (tx_a, mut rx_a) = watch::channel::>(None); let (tx_b, mut rx_b) = watch::channel::>(None); store.update(cx, |store, _| { store .external_agents .get_mut(&AgentId::new("agent-a")) .expect("agent-a should be registered") .server .set_new_version_available_tx(tx_a); store .external_agents .get_mut(&AgentId::new("agent-b")) .expect("agent-b should be registered") .server .set_new_version_available_tx(tx_b); }); // Update only agent-a to a new version; agent-b stays the same. registry.update(cx, |store, cx| { store.set_agents( vec![ make_npx_agent("agent-a", "2.0.0"), make_npx_agent("agent-b", "3.0.0"), ], cx, ); }); cx.run_until_parked(); // agent-a should have received a notification. assert_eq!(rx_a.borrow().as_deref(), Some("2.0.0")); // agent-b should NOT have received a notification. assert_eq!(rx_b.borrow().as_deref(), None); // agent-b's tx should have been transferred. store.update(cx, |store, _| { assert!( store .external_agents .get_mut(&AgentId::new("agent-b")) .expect("agent-b should be registered") .server .take_new_version_available_tx() .is_some(), "agent-b tx should have been transferred" ); }); } }