agent_server_store.rs

   1use std::{
   2    any::Any,
   3    path::{Path, PathBuf},
   4    sync::Arc,
   5    time::Duration,
   6};
   7
   8use anyhow::{Context as _, Result, bail};
   9use collections::HashMap;
  10use fs::Fs;
  11use gpui::{AsyncApp, Context, Entity, EventEmitter, SharedString, Subscription, Task};
  12use http_client::{HttpClient, github::AssetKind};
  13use node_runtime::NodeRuntime;
  14use percent_encoding::percent_decode_str;
  15use remote::RemoteClient;
  16use rpc::{
  17    AnyProtoClient, TypedEnvelope,
  18    proto::{self, ExternalExtensionAgent},
  19};
  20use schemars::JsonSchema;
  21use serde::{Deserialize, Serialize};
  22use settings::{RegisterSetting, SettingsStore};
  23use sha2::{Digest, Sha256};
  24use url::Url;
  25use util::{ResultExt as _, debug_panic};
  26
  27use crate::ProjectEnvironment;
  28use crate::agent_registry_store::{AgentRegistryStore, RegistryAgent, RegistryTargetConfig};
  29
  30use crate::worktree_store::WorktreeStore;
  31
  32#[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema)]
  33pub struct AgentServerCommand {
  34    #[serde(rename = "command")]
  35    pub path: PathBuf,
  36    #[serde(default)]
  37    pub args: Vec<String>,
  38    pub env: Option<HashMap<String, String>>,
  39}
  40
  41impl std::fmt::Debug for AgentServerCommand {
  42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  43        let filtered_env = self.env.as_ref().map(|env| {
  44            env.iter()
  45                .map(|(k, v)| {
  46                    (
  47                        k,
  48                        if util::redact::should_redact(k) {
  49                            "[REDACTED]"
  50                        } else {
  51                            v
  52                        },
  53                    )
  54                })
  55                .collect::<Vec<_>>()
  56        });
  57
  58        f.debug_struct("AgentServerCommand")
  59            .field("path", &self.path)
  60            .field("args", &self.args)
  61            .field("env", &filtered_env)
  62            .finish()
  63    }
  64}
  65
  66#[derive(
  67    Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, JsonSchema,
  68)]
  69#[serde(transparent)]
  70pub struct AgentId(pub SharedString);
  71
  72impl AgentId {
  73    pub fn new(id: impl Into<SharedString>) -> Self {
  74        AgentId(id.into())
  75    }
  76}
  77
  78impl std::fmt::Display for AgentId {
  79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  80        write!(f, "{}", self.0)
  81    }
  82}
  83
  84impl From<&'static str> for AgentId {
  85    fn from(value: &'static str) -> Self {
  86        AgentId(value.into())
  87    }
  88}
  89
  90impl From<AgentId> for SharedString {
  91    fn from(value: AgentId) -> Self {
  92        value.0
  93    }
  94}
  95
  96impl AsRef<str> for AgentId {
  97    fn as_ref(&self) -> &str {
  98        &self.0
  99    }
 100}
 101
 102impl std::borrow::Borrow<str> for AgentId {
 103    fn borrow(&self) -> &str {
 104        &self.0
 105    }
 106}
 107
 108#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 109pub enum ExternalAgentSource {
 110    #[default]
 111    Custom,
 112    Extension,
 113    Registry,
 114}
 115
 116pub trait ExternalAgentServer {
 117    fn get_command(
 118        &self,
 119        extra_args: Vec<String>,
 120        extra_env: HashMap<String, String>,
 121        cx: &mut AsyncApp,
 122    ) -> Task<Result<AgentServerCommand>>;
 123
 124    fn version(&self) -> Option<&SharedString> {
 125        None
 126    }
 127
 128    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
 129        None
 130    }
 131
 132    fn set_new_version_available_tx(&mut self, _tx: watch::Sender<Option<String>>) {}
 133
 134    fn as_any(&self) -> &dyn Any;
 135    fn as_any_mut(&mut self) -> &mut dyn Any;
 136}
 137
 138struct ExtensionAgentEntry {
 139    agent_name: Arc<str>,
 140    extension_id: String,
 141    targets: HashMap<String, extension::TargetConfig>,
 142    env: HashMap<String, String>,
 143    icon_path: Option<String>,
 144    display_name: Option<SharedString>,
 145    version: Option<SharedString>,
 146}
 147
 148enum AgentServerStoreState {
 149    Local {
 150        node_runtime: NodeRuntime,
 151        fs: Arc<dyn Fs>,
 152        project_environment: Entity<ProjectEnvironment>,
 153        downstream_client: Option<(u64, AnyProtoClient)>,
 154        settings: Option<AllAgentServersSettings>,
 155        http_client: Arc<dyn HttpClient>,
 156        extension_agents: Vec<ExtensionAgentEntry>,
 157        _subscriptions: Vec<Subscription>,
 158    },
 159    Remote {
 160        project_id: u64,
 161        upstream_client: Entity<RemoteClient>,
 162        worktree_store: Entity<WorktreeStore>,
 163    },
 164    Collab,
 165}
 166
 167pub struct ExternalAgentEntry {
 168    server: Box<dyn ExternalAgentServer>,
 169    icon: Option<SharedString>,
 170    display_name: Option<SharedString>,
 171    pub source: ExternalAgentSource,
 172}
 173
 174impl ExternalAgentEntry {
 175    pub fn new(
 176        server: Box<dyn ExternalAgentServer>,
 177        source: ExternalAgentSource,
 178        icon: Option<SharedString>,
 179        display_name: Option<SharedString>,
 180    ) -> Self {
 181        Self {
 182            server,
 183            icon,
 184            display_name,
 185            source,
 186        }
 187    }
 188}
 189
 190pub struct AgentServerStore {
 191    state: AgentServerStoreState,
 192    pub external_agents: HashMap<AgentId, ExternalAgentEntry>,
 193}
 194
 195pub struct AgentServersUpdated;
 196
 197impl EventEmitter<AgentServersUpdated> for AgentServerStore {}
 198
 199impl AgentServerStore {
 200    /// Synchronizes extension-provided agent servers with the store.
 201    pub fn sync_extension_agents<'a, I>(
 202        &mut self,
 203        manifests: I,
 204        extensions_dir: PathBuf,
 205        cx: &mut Context<Self>,
 206    ) where
 207        I: IntoIterator<Item = (&'a str, &'a extension::ExtensionManifest)>,
 208    {
 209        // Collect manifests first so we can iterate twice
 210        let manifests: Vec<_> = manifests.into_iter().collect();
 211
 212        // Remove all extension-provided agents
 213        // (They will be re-added below if they're in the currently installed extensions)
 214        self.external_agents
 215            .retain(|_, entry| entry.source != ExternalAgentSource::Extension);
 216
 217        // Insert agent servers from extension manifests
 218        match &mut self.state {
 219            AgentServerStoreState::Local {
 220                extension_agents, ..
 221            } => {
 222                extension_agents.clear();
 223                for (ext_id, manifest) in manifests {
 224                    for (agent_name, agent_entry) in &manifest.agent_servers {
 225                        let display_name = SharedString::from(agent_entry.name.clone());
 226                        let icon_path = agent_entry.icon.as_ref().and_then(|icon| {
 227                            resolve_extension_icon_path(&extensions_dir, ext_id, icon)
 228                        });
 229
 230                        extension_agents.push(ExtensionAgentEntry {
 231                            agent_name: agent_name.clone(),
 232                            extension_id: ext_id.to_owned(),
 233                            targets: agent_entry.targets.clone(),
 234                            env: agent_entry.env.clone(),
 235                            icon_path,
 236                            display_name: Some(display_name),
 237                            version: Some(SharedString::from(manifest.version.clone())),
 238                        });
 239                    }
 240                }
 241                self.reregister_agents(cx);
 242            }
 243            AgentServerStoreState::Remote {
 244                project_id,
 245                upstream_client,
 246                worktree_store,
 247            } => {
 248                let mut agents = vec![];
 249                for (ext_id, manifest) in manifests {
 250                    for (agent_name, agent_entry) in &manifest.agent_servers {
 251                        let display_name = SharedString::from(agent_entry.name.clone());
 252                        let icon_path = agent_entry.icon.as_ref().and_then(|icon| {
 253                            resolve_extension_icon_path(&extensions_dir, ext_id, icon)
 254                        });
 255                        let icon_shared = icon_path
 256                            .as_ref()
 257                            .map(|path| SharedString::from(path.clone()));
 258                        let icon = icon_path;
 259                        let agent_server_name = AgentId(agent_name.clone().into());
 260                        self.external_agents
 261                            .entry(agent_server_name.clone())
 262                            .and_modify(|entry| {
 263                                entry.icon = icon_shared.clone();
 264                                entry.display_name = Some(display_name.clone());
 265                                entry.source = ExternalAgentSource::Extension;
 266                            })
 267                            .or_insert_with(|| {
 268                                ExternalAgentEntry::new(
 269                                    Box::new(RemoteExternalAgentServer {
 270                                        project_id: *project_id,
 271                                        upstream_client: upstream_client.clone(),
 272                                        worktree_store: worktree_store.clone(),
 273                                        name: agent_server_name.clone(),
 274                                        new_version_available_tx: None,
 275                                    })
 276                                        as Box<dyn ExternalAgentServer>,
 277                                    ExternalAgentSource::Extension,
 278                                    icon_shared.clone(),
 279                                    Some(display_name.clone()),
 280                                )
 281                            });
 282
 283                        agents.push(ExternalExtensionAgent {
 284                            name: agent_name.to_string(),
 285                            icon_path: icon,
 286                            extension_id: ext_id.to_string(),
 287                            targets: agent_entry
 288                                .targets
 289                                .iter()
 290                                .map(|(k, v)| (k.clone(), v.to_proto()))
 291                                .collect(),
 292                            env: agent_entry
 293                                .env
 294                                .iter()
 295                                .map(|(k, v)| (k.clone(), v.clone()))
 296                                .collect(),
 297                            version: Some(manifest.version.to_string()),
 298                        });
 299                    }
 300                }
 301                upstream_client
 302                    .read(cx)
 303                    .proto_client()
 304                    .send(proto::ExternalExtensionAgentsUpdated {
 305                        project_id: *project_id,
 306                        agents,
 307                    })
 308                    .log_err();
 309            }
 310            AgentServerStoreState::Collab => {
 311                // Do nothing
 312            }
 313        }
 314
 315        cx.emit(AgentServersUpdated);
 316    }
 317
 318    pub fn agent_icon(&self, id: &AgentId) -> Option<SharedString> {
 319        self.external_agents
 320            .get(id)
 321            .and_then(|entry| entry.icon.clone())
 322    }
 323
 324    pub fn agent_source(&self, name: &AgentId) -> Option<ExternalAgentSource> {
 325        self.external_agents.get(name).map(|entry| entry.source)
 326    }
 327}
 328
 329/// Safely resolves an extension icon path, ensuring it stays within the extension directory.
 330/// Returns `None` if the path would escape the extension directory (path traversal attack).
 331pub fn resolve_extension_icon_path(
 332    extensions_dir: &Path,
 333    extension_id: &str,
 334    icon_relative_path: &str,
 335) -> Option<String> {
 336    let extension_root = extensions_dir.join(extension_id);
 337    let icon_path = extension_root.join(icon_relative_path);
 338
 339    // Canonicalize both paths to resolve symlinks and normalize the paths.
 340    // For the extension root, we need to handle the case where it might be a symlink
 341    // (common for dev extensions).
 342    let canonical_extension_root = extension_root.canonicalize().unwrap_or(extension_root);
 343    let canonical_icon_path = match icon_path.canonicalize() {
 344        Ok(path) => path,
 345        Err(err) => {
 346            log::warn!(
 347                "Failed to canonicalize icon path for extension '{}': {} (path: {})",
 348                extension_id,
 349                err,
 350                icon_relative_path
 351            );
 352            return None;
 353        }
 354    };
 355
 356    // Verify the resolved icon path is within the extension directory
 357    if canonical_icon_path.starts_with(&canonical_extension_root) {
 358        Some(canonical_icon_path.to_string_lossy().to_string())
 359    } else {
 360        log::warn!(
 361            "Icon path '{}' for extension '{}' escapes extension directory, ignoring for security",
 362            icon_relative_path,
 363            extension_id
 364        );
 365        None
 366    }
 367}
 368
 369impl AgentServerStore {
 370    pub fn agent_display_name(&self, name: &AgentId) -> Option<SharedString> {
 371        self.external_agents
 372            .get(name)
 373            .and_then(|entry| entry.display_name.clone())
 374    }
 375
 376    pub fn init_remote(session: &AnyProtoClient) {
 377        session.add_entity_message_handler(Self::handle_external_agents_updated);
 378        session.add_entity_message_handler(Self::handle_new_version_available);
 379    }
 380
 381    pub fn init_headless(session: &AnyProtoClient) {
 382        session.add_entity_message_handler(Self::handle_external_extension_agents_updated);
 383        session.add_entity_request_handler(Self::handle_get_agent_server_command);
 384    }
 385
 386    fn agent_servers_settings_changed(&mut self, cx: &mut Context<Self>) {
 387        let AgentServerStoreState::Local {
 388            settings: old_settings,
 389            ..
 390        } = &mut self.state
 391        else {
 392            debug_panic!(
 393                "should not be subscribed to agent server settings changes in non-local project"
 394            );
 395            return;
 396        };
 397
 398        let new_settings = cx
 399            .global::<SettingsStore>()
 400            .get::<AllAgentServersSettings>(None)
 401            .clone();
 402        if Some(&new_settings) == old_settings.as_ref() {
 403            return;
 404        }
 405
 406        self.reregister_agents(cx);
 407    }
 408
 409    fn reregister_agents(&mut self, cx: &mut Context<Self>) {
 410        let AgentServerStoreState::Local {
 411            node_runtime,
 412            fs,
 413            project_environment,
 414            downstream_client,
 415            settings: old_settings,
 416            http_client,
 417            extension_agents,
 418            ..
 419        } = &mut self.state
 420        else {
 421            debug_panic!("Non-local projects should never attempt to reregister. This is a bug!");
 422
 423            return;
 424        };
 425
 426        let new_settings = cx
 427            .global::<SettingsStore>()
 428            .get::<AllAgentServersSettings>(None)
 429            .clone();
 430
 431        // If we don't have agents from the registry loaded yet, trigger a
 432        // refresh, which will cause this function to be called again
 433        let registry_store = AgentRegistryStore::try_global(cx);
 434        if new_settings.has_registry_agents()
 435            && let Some(registry) = registry_store.as_ref()
 436        {
 437            registry.update(cx, |registry, cx| registry.refresh_if_stale(cx));
 438        }
 439
 440        let registry_agents_by_id = registry_store
 441            .as_ref()
 442            .map(|store| {
 443                store
 444                    .read(cx)
 445                    .agents()
 446                    .iter()
 447                    .cloned()
 448                    .map(|agent| (agent.id().to_string(), agent))
 449                    .collect::<HashMap<_, _>>()
 450            })
 451            .unwrap_or_default();
 452
 453        // Drain the existing versioned agents, extracting reconnect state
 454        // from any active connection so we can preserve it or trigger a
 455        // reconnect when the version changes.
 456        let mut old_versioned_agents: HashMap<
 457            AgentId,
 458            (SharedString, watch::Sender<Option<String>>),
 459        > = HashMap::default();
 460        for (name, mut entry) in self.external_agents.drain() {
 461            if let Some(version) = entry.server.version().cloned() {
 462                if let Some(tx) = entry.server.take_new_version_available_tx() {
 463                    old_versioned_agents.insert(name, (version, tx));
 464                }
 465            }
 466        }
 467
 468        // Insert extension agents before custom/registry so registry entries override extensions.
 469        for entry in extension_agents.iter() {
 470            let name = AgentId(entry.agent_name.clone().into());
 471            let mut env = entry.env.clone();
 472            if let Some(settings_env) =
 473                new_settings
 474                    .get(entry.agent_name.as_ref())
 475                    .and_then(|settings| match settings {
 476                        CustomAgentServerSettings::Extension { env, .. } => Some(env.clone()),
 477                        _ => None,
 478                    })
 479            {
 480                env.extend(settings_env);
 481            }
 482            let icon = entry
 483                .icon_path
 484                .as_ref()
 485                .map(|path| SharedString::from(path.clone()));
 486
 487            self.external_agents.insert(
 488                name.clone(),
 489                ExternalAgentEntry::new(
 490                    Box::new(LocalExtensionArchiveAgent {
 491                        fs: fs.clone(),
 492                        http_client: http_client.clone(),
 493                        node_runtime: node_runtime.clone(),
 494                        project_environment: project_environment.clone(),
 495                        extension_id: Arc::from(&*entry.extension_id),
 496                        targets: entry.targets.clone(),
 497                        env,
 498                        agent_id: entry.agent_name.clone(),
 499                        version: entry.version.clone(),
 500                        new_version_available_tx: None,
 501                    }) as Box<dyn ExternalAgentServer>,
 502                    ExternalAgentSource::Extension,
 503                    icon,
 504                    entry.display_name.clone(),
 505                ),
 506            );
 507        }
 508
 509        for (name, settings) in new_settings.iter() {
 510            match settings {
 511                CustomAgentServerSettings::Custom { command, .. } => {
 512                    let agent_name = AgentId(name.clone().into());
 513                    self.external_agents.insert(
 514                        agent_name.clone(),
 515                        ExternalAgentEntry::new(
 516                            Box::new(LocalCustomAgent {
 517                                command: command.clone(),
 518                                project_environment: project_environment.clone(),
 519                            }) as Box<dyn ExternalAgentServer>,
 520                            ExternalAgentSource::Custom,
 521                            None,
 522                            None,
 523                        ),
 524                    );
 525                }
 526                CustomAgentServerSettings::Registry { env, .. } => {
 527                    let Some(agent) = registry_agents_by_id.get(name) else {
 528                        if registry_store.is_some() {
 529                            log::debug!("Registry agent '{}' not found in ACP registry", name);
 530                        }
 531                        continue;
 532                    };
 533
 534                    let agent_name = AgentId(name.clone().into());
 535                    match agent {
 536                        RegistryAgent::Binary(agent) => {
 537                            if !agent.supports_current_platform {
 538                                log::warn!(
 539                                    "Registry agent '{}' has no compatible binary for this platform",
 540                                    name
 541                                );
 542                                continue;
 543                            }
 544
 545                            self.external_agents.insert(
 546                                agent_name.clone(),
 547                                ExternalAgentEntry::new(
 548                                    Box::new(LocalRegistryArchiveAgent {
 549                                        fs: fs.clone(),
 550                                        http_client: http_client.clone(),
 551                                        node_runtime: node_runtime.clone(),
 552                                        project_environment: project_environment.clone(),
 553                                        registry_id: Arc::from(name.as_str()),
 554                                        version: agent.metadata.version.clone(),
 555                                        targets: agent.targets.clone(),
 556                                        env: env.clone(),
 557                                        new_version_available_tx: None,
 558                                    })
 559                                        as Box<dyn ExternalAgentServer>,
 560                                    ExternalAgentSource::Registry,
 561                                    agent.metadata.icon_path.clone(),
 562                                    Some(agent.metadata.name.clone()),
 563                                ),
 564                            );
 565                        }
 566                        RegistryAgent::Npx(agent) => {
 567                            self.external_agents.insert(
 568                                agent_name.clone(),
 569                                ExternalAgentEntry::new(
 570                                    Box::new(LocalRegistryNpxAgent {
 571                                        fs: fs.clone(),
 572                                        node_runtime: node_runtime.clone(),
 573                                        project_environment: project_environment.clone(),
 574                                        version: agent.metadata.version.clone(),
 575                                        package: agent.package.clone(),
 576                                        args: agent.args.clone(),
 577                                        distribution_env: agent.env.clone(),
 578                                        settings_env: env.clone(),
 579                                        new_version_available_tx: None,
 580                                    })
 581                                        as Box<dyn ExternalAgentServer>,
 582                                    ExternalAgentSource::Registry,
 583                                    agent.metadata.icon_path.clone(),
 584                                    Some(agent.metadata.name.clone()),
 585                                ),
 586                            );
 587                        }
 588                    }
 589                }
 590                CustomAgentServerSettings::Extension { .. } => {}
 591            }
 592        }
 593
 594        // For each rebuilt versioned agent, compare the version. If it
 595        // changed, notify the active connection to reconnect. Otherwise,
 596        // transfer the channel to the new entry so future updates can use it.
 597        for (name, entry) in &mut self.external_agents {
 598            let Some((old_version, mut tx)) = old_versioned_agents.remove(name) else {
 599                continue;
 600            };
 601            let Some(new_version) = entry.server.version() else {
 602                continue;
 603            };
 604
 605            if new_version != &old_version {
 606                tx.send(Some(new_version.to_string())).ok();
 607            } else {
 608                entry.server.set_new_version_available_tx(tx);
 609            }
 610        }
 611
 612        *old_settings = Some(new_settings);
 613
 614        if let Some((project_id, downstream_client)) = downstream_client {
 615            downstream_client
 616                .send(proto::ExternalAgentsUpdated {
 617                    project_id: *project_id,
 618                    names: self
 619                        .external_agents
 620                        .keys()
 621                        .map(|name| name.to_string())
 622                        .collect(),
 623                })
 624                .log_err();
 625        }
 626        cx.emit(AgentServersUpdated);
 627    }
 628
 629    pub fn node_runtime(&self) -> Option<NodeRuntime> {
 630        match &self.state {
 631            AgentServerStoreState::Local { node_runtime, .. } => Some(node_runtime.clone()),
 632            _ => None,
 633        }
 634    }
 635
 636    pub fn local(
 637        node_runtime: NodeRuntime,
 638        fs: Arc<dyn Fs>,
 639        project_environment: Entity<ProjectEnvironment>,
 640        http_client: Arc<dyn HttpClient>,
 641        cx: &mut Context<Self>,
 642    ) -> Self {
 643        let mut subscriptions = vec![cx.observe_global::<SettingsStore>(|this, cx| {
 644            this.agent_servers_settings_changed(cx);
 645        })];
 646        if let Some(registry_store) = AgentRegistryStore::try_global(cx) {
 647            subscriptions.push(cx.observe(&registry_store, |this, _, cx| {
 648                this.reregister_agents(cx);
 649            }));
 650        }
 651        let mut this = Self {
 652            state: AgentServerStoreState::Local {
 653                node_runtime,
 654                fs,
 655                project_environment,
 656                http_client,
 657                downstream_client: None,
 658                settings: None,
 659                extension_agents: vec![],
 660                _subscriptions: subscriptions,
 661            },
 662            external_agents: HashMap::default(),
 663        };
 664        if let Some(_events) = extension::ExtensionEvents::try_global(cx) {}
 665        this.agent_servers_settings_changed(cx);
 666        this
 667    }
 668
 669    pub(crate) fn remote(
 670        project_id: u64,
 671        upstream_client: Entity<RemoteClient>,
 672        worktree_store: Entity<WorktreeStore>,
 673    ) -> Self {
 674        Self {
 675            state: AgentServerStoreState::Remote {
 676                project_id,
 677                upstream_client,
 678                worktree_store,
 679            },
 680            external_agents: HashMap::default(),
 681        }
 682    }
 683
 684    pub fn collab() -> Self {
 685        Self {
 686            state: AgentServerStoreState::Collab,
 687            external_agents: HashMap::default(),
 688        }
 689    }
 690
 691    pub fn shared(&mut self, project_id: u64, client: AnyProtoClient, cx: &mut Context<Self>) {
 692        match &mut self.state {
 693            AgentServerStoreState::Local {
 694                downstream_client, ..
 695            } => {
 696                *downstream_client = Some((project_id, client.clone()));
 697                // Send the current list of external agents downstream, but only after a delay,
 698                // to avoid having the message arrive before the downstream project's agent server store
 699                // sets up its handlers.
 700                cx.spawn(async move |this, cx| {
 701                    cx.background_executor().timer(Duration::from_secs(1)).await;
 702                    let names = this.update(cx, |this, _| {
 703                        this.external_agents()
 704                            .map(|name| name.to_string())
 705                            .collect()
 706                    })?;
 707                    client
 708                        .send(proto::ExternalAgentsUpdated { project_id, names })
 709                        .log_err();
 710                    anyhow::Ok(())
 711                })
 712                .detach();
 713            }
 714            AgentServerStoreState::Remote { .. } => {
 715                debug_panic!(
 716                    "external agents over collab not implemented, remote project should not be shared"
 717                );
 718            }
 719            AgentServerStoreState::Collab => {
 720                debug_panic!("external agents over collab not implemented, should not be shared");
 721            }
 722        }
 723    }
 724
 725    pub fn get_external_agent(
 726        &mut self,
 727        name: &AgentId,
 728    ) -> Option<&mut (dyn ExternalAgentServer + 'static)> {
 729        self.external_agents
 730            .get_mut(name)
 731            .map(|entry| entry.server.as_mut())
 732    }
 733
 734    pub fn no_browser(&self) -> bool {
 735        match &self.state {
 736            AgentServerStoreState::Local {
 737                downstream_client, ..
 738            } => downstream_client
 739                .as_ref()
 740                .is_some_and(|(_, client)| !client.has_wsl_interop()),
 741            _ => false,
 742        }
 743    }
 744
 745    pub fn has_external_agents(&self) -> bool {
 746        !self.external_agents.is_empty()
 747    }
 748
 749    pub fn external_agents(&self) -> impl Iterator<Item = &AgentId> {
 750        self.external_agents.keys()
 751    }
 752
 753    async fn handle_get_agent_server_command(
 754        this: Entity<Self>,
 755        envelope: TypedEnvelope<proto::GetAgentServerCommand>,
 756        mut cx: AsyncApp,
 757    ) -> Result<proto::AgentServerCommand> {
 758        let command = this
 759            .update(&mut cx, |this, cx| {
 760                let AgentServerStoreState::Local {
 761                    downstream_client, ..
 762                } = &this.state
 763                else {
 764                    debug_panic!("should not receive GetAgentServerCommand in a non-local project");
 765                    bail!("unexpected GetAgentServerCommand request in a non-local project");
 766                };
 767                let no_browser = this.no_browser();
 768                let agent = this
 769                    .external_agents
 770                    .get_mut(&*envelope.payload.name)
 771                    .map(|entry| entry.server.as_mut())
 772                    .with_context(|| format!("agent `{}` not found", envelope.payload.name))?;
 773                let new_version_available_tx =
 774                    downstream_client
 775                        .clone()
 776                        .map(|(project_id, downstream_client)| {
 777                            let (new_version_available_tx, mut new_version_available_rx) =
 778                                watch::channel(None);
 779                            cx.spawn({
 780                                let name = envelope.payload.name.clone();
 781                                async move |_, _| {
 782                                    if let Some(version) =
 783                                        new_version_available_rx.recv().await.ok().flatten()
 784                                    {
 785                                        downstream_client.send(
 786                                            proto::NewExternalAgentVersionAvailable {
 787                                                project_id,
 788                                                name: name.clone(),
 789                                                version,
 790                                            },
 791                                        )?;
 792                                    }
 793                                    anyhow::Ok(())
 794                                }
 795                            })
 796                            .detach_and_log_err(cx);
 797                            new_version_available_tx
 798                        });
 799                let mut extra_env = HashMap::default();
 800                if no_browser {
 801                    extra_env.insert("NO_BROWSER".to_owned(), "1".to_owned());
 802                }
 803                if let Some(new_version_available_tx) = new_version_available_tx {
 804                    agent.set_new_version_available_tx(new_version_available_tx);
 805                }
 806                anyhow::Ok(agent.get_command(vec![], extra_env, &mut cx.to_async()))
 807            })?
 808            .await?;
 809        Ok(proto::AgentServerCommand {
 810            path: command.path.to_string_lossy().into_owned(),
 811            args: command.args,
 812            env: command
 813                .env
 814                .map(|env| env.into_iter().collect())
 815                .unwrap_or_default(),
 816            root_dir: envelope
 817                .payload
 818                .root_dir
 819                .unwrap_or_else(|| paths::home_dir().to_string_lossy().to_string()),
 820            login: None,
 821        })
 822    }
 823
 824    async fn handle_external_agents_updated(
 825        this: Entity<Self>,
 826        envelope: TypedEnvelope<proto::ExternalAgentsUpdated>,
 827        mut cx: AsyncApp,
 828    ) -> Result<()> {
 829        this.update(&mut cx, |this, cx| {
 830            let AgentServerStoreState::Remote {
 831                project_id,
 832                upstream_client,
 833                worktree_store,
 834            } = &this.state
 835            else {
 836                debug_panic!(
 837                    "handle_external_agents_updated should not be called for a non-remote project"
 838                );
 839                bail!("unexpected ExternalAgentsUpdated message")
 840            };
 841
 842            let mut previous_entries = std::mem::take(&mut this.external_agents);
 843            let mut new_version_available_txs = HashMap::default();
 844            let mut metadata = HashMap::default();
 845
 846            for (name, mut entry) in previous_entries.drain() {
 847                if let Some(tx) = entry.server.take_new_version_available_tx() {
 848                    new_version_available_txs.insert(name.clone(), tx);
 849                }
 850
 851                metadata.insert(name, (entry.icon, entry.display_name, entry.source));
 852            }
 853
 854            this.external_agents = envelope
 855                .payload
 856                .names
 857                .into_iter()
 858                .map(|name| {
 859                    let agent_id = AgentId(name.into());
 860                    let (icon, display_name, source) = metadata
 861                        .remove(&agent_id)
 862                        .or_else(|| {
 863                            AgentRegistryStore::try_global(cx)
 864                                .and_then(|store| store.read(cx).agent(&agent_id))
 865                                .map(|s| {
 866                                    (
 867                                        s.icon_path().cloned(),
 868                                        Some(s.name().clone()),
 869                                        ExternalAgentSource::Registry,
 870                                    )
 871                                })
 872                        })
 873                        .unwrap_or((None, None, ExternalAgentSource::default()));
 874                    let agent = RemoteExternalAgentServer {
 875                        project_id: *project_id,
 876                        upstream_client: upstream_client.clone(),
 877                        worktree_store: worktree_store.clone(),
 878                        name: agent_id.clone(),
 879                        new_version_available_tx: new_version_available_txs.remove(&agent_id),
 880                    };
 881                    (
 882                        agent_id,
 883                        ExternalAgentEntry::new(
 884                            Box::new(agent) as Box<dyn ExternalAgentServer>,
 885                            source,
 886                            icon,
 887                            display_name,
 888                        ),
 889                    )
 890                })
 891                .collect();
 892            cx.emit(AgentServersUpdated);
 893            Ok(())
 894        })
 895    }
 896
 897    async fn handle_external_extension_agents_updated(
 898        this: Entity<Self>,
 899        envelope: TypedEnvelope<proto::ExternalExtensionAgentsUpdated>,
 900        mut cx: AsyncApp,
 901    ) -> Result<()> {
 902        this.update(&mut cx, |this, cx| {
 903            let AgentServerStoreState::Local {
 904                extension_agents, ..
 905            } = &mut this.state
 906            else {
 907                panic!(
 908                    "handle_external_extension_agents_updated \
 909                    should not be called for a non-remote project"
 910                );
 911            };
 912
 913            extension_agents.clear();
 914            for ExternalExtensionAgent {
 915                name,
 916                icon_path,
 917                extension_id,
 918                targets,
 919                env,
 920                version,
 921            } in envelope.payload.agents
 922            {
 923                extension_agents.push(ExtensionAgentEntry {
 924                    agent_name: Arc::from(&*name),
 925                    extension_id,
 926                    targets: targets
 927                        .into_iter()
 928                        .map(|(k, v)| (k, extension::TargetConfig::from_proto(v)))
 929                        .collect(),
 930                    env: env.into_iter().collect(),
 931                    icon_path,
 932                    display_name: None,
 933                    version: version.map(SharedString::from),
 934                });
 935            }
 936
 937            this.reregister_agents(cx);
 938            cx.emit(AgentServersUpdated);
 939            Ok(())
 940        })
 941    }
 942
 943    async fn handle_new_version_available(
 944        this: Entity<Self>,
 945        envelope: TypedEnvelope<proto::NewExternalAgentVersionAvailable>,
 946        mut cx: AsyncApp,
 947    ) -> Result<()> {
 948        this.update(&mut cx, |this, _| {
 949            if let Some(entry) = this.external_agents.get_mut(&*envelope.payload.name)
 950                && let Some(mut tx) = entry.server.take_new_version_available_tx()
 951            {
 952                tx.send(Some(envelope.payload.version)).ok();
 953                entry.server.set_new_version_available_tx(tx);
 954            }
 955        });
 956        Ok(())
 957    }
 958
 959    pub fn get_extension_id_for_agent(&self, name: &AgentId) -> Option<Arc<str>> {
 960        self.external_agents.get(name).and_then(|entry| {
 961            entry
 962                .server
 963                .as_any()
 964                .downcast_ref::<LocalExtensionArchiveAgent>()
 965                .map(|ext_agent| ext_agent.extension_id.clone())
 966        })
 967    }
 968}
 969
 970struct RemoteExternalAgentServer {
 971    project_id: u64,
 972    upstream_client: Entity<RemoteClient>,
 973    worktree_store: Entity<WorktreeStore>,
 974    name: AgentId,
 975    new_version_available_tx: Option<watch::Sender<Option<String>>>,
 976}
 977
 978impl ExternalAgentServer for RemoteExternalAgentServer {
 979    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
 980        self.new_version_available_tx.take()
 981    }
 982
 983    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
 984        self.new_version_available_tx = Some(tx);
 985    }
 986
 987    fn get_command(
 988        &self,
 989        extra_args: Vec<String>,
 990        extra_env: HashMap<String, String>,
 991        cx: &mut AsyncApp,
 992    ) -> Task<Result<AgentServerCommand>> {
 993        let project_id = self.project_id;
 994        let name = self.name.to_string();
 995        let upstream_client = self.upstream_client.downgrade();
 996        let worktree_store = self.worktree_store.clone();
 997        cx.spawn(async move |cx| {
 998            let root_dir = worktree_store.read_with(cx, |worktree_store, cx| {
 999                crate::Project::default_visible_worktree_paths(worktree_store, cx)
1000                    .into_iter()
1001                    .next()
1002                    .map(|path| path.display().to_string())
1003            });
1004
1005            let mut response = upstream_client
1006                .update(cx, |upstream_client, _| {
1007                    upstream_client
1008                        .proto_client()
1009                        .request(proto::GetAgentServerCommand {
1010                            project_id,
1011                            name,
1012                            root_dir,
1013                        })
1014                })?
1015                .await?;
1016            response.args.extend(extra_args);
1017            response.env.extend(extra_env);
1018
1019            Ok(AgentServerCommand {
1020                path: response.path.into(),
1021                args: response.args,
1022                env: Some(response.env.into_iter().collect()),
1023            })
1024        })
1025    }
1026
1027    fn as_any(&self) -> &dyn Any {
1028        self
1029    }
1030
1031    fn as_any_mut(&mut self) -> &mut dyn Any {
1032        self
1033    }
1034}
1035
1036fn asset_kind_for_archive_url(archive_url: &str) -> Result<AssetKind> {
1037    let archive_path = Url::parse(archive_url)
1038        .ok()
1039        .map(|url| url.path().to_string())
1040        .unwrap_or_else(|| archive_url.to_string());
1041
1042    if archive_path.ends_with(".zip") {
1043        Ok(AssetKind::Zip)
1044    } else if archive_path.ends_with(".tar.gz") || archive_path.ends_with(".tgz") {
1045        Ok(AssetKind::TarGz)
1046    } else if archive_path.ends_with(".tar.bz2") || archive_path.ends_with(".tbz2") {
1047        Ok(AssetKind::TarBz2)
1048    } else {
1049        bail!("unsupported archive type in URL: {archive_url}");
1050    }
1051}
1052
1053struct GithubReleaseArchive {
1054    repo_name_with_owner: String,
1055    tag: String,
1056    asset_name: String,
1057}
1058
1059fn github_release_archive_from_url(archive_url: &str) -> Option<GithubReleaseArchive> {
1060    fn decode_path_segment(segment: &str) -> Option<String> {
1061        percent_decode_str(segment)
1062            .decode_utf8()
1063            .ok()
1064            .map(|segment| segment.into_owned())
1065    }
1066
1067    let url = Url::parse(archive_url).ok()?;
1068    if url.scheme() != "https" || url.host_str()? != "github.com" {
1069        return None;
1070    }
1071
1072    let segments = url.path_segments()?.collect::<Vec<_>>();
1073    if segments.len() < 6 || segments[2] != "releases" || segments[3] != "download" {
1074        return None;
1075    }
1076
1077    Some(GithubReleaseArchive {
1078        repo_name_with_owner: format!("{}/{}", segments[0], segments[1]),
1079        tag: decode_path_segment(segments[4])?,
1080        asset_name: segments[5..]
1081            .iter()
1082            .map(|segment| decode_path_segment(segment))
1083            .collect::<Option<Vec<_>>>()?
1084            .join("/"),
1085    })
1086}
1087
1088fn sanitized_version_component(version: &str) -> String {
1089    let sanitized = version
1090        .chars()
1091        .map(|character| match character {
1092            'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '_' | '-' => character,
1093            _ => '-',
1094        })
1095        .collect::<String>();
1096
1097    if sanitized.is_empty() {
1098        "unknown".to_string()
1099    } else {
1100        sanitized
1101    }
1102}
1103
1104fn versioned_archive_cache_dir(
1105    base_dir: &Path,
1106    version: Option<&str>,
1107    archive_url: &str,
1108) -> PathBuf {
1109    let version = version.unwrap_or_default();
1110    let sanitized_version = sanitized_version_component(version);
1111
1112    let mut version_hasher = Sha256::new();
1113    version_hasher.update(version.as_bytes());
1114    let version_hash = format!("{:x}", version_hasher.finalize());
1115
1116    let mut url_hasher = Sha256::new();
1117    url_hasher.update(archive_url.as_bytes());
1118    let url_hash = format!("{:x}", url_hasher.finalize());
1119
1120    base_dir.join(format!(
1121        "v_{sanitized_version}_{}_{}",
1122        &version_hash[..16],
1123        &url_hash[..16],
1124    ))
1125}
1126
1127pub struct LocalExtensionArchiveAgent {
1128    pub fs: Arc<dyn Fs>,
1129    pub http_client: Arc<dyn HttpClient>,
1130    pub node_runtime: NodeRuntime,
1131    pub project_environment: Entity<ProjectEnvironment>,
1132    pub extension_id: Arc<str>,
1133    pub agent_id: Arc<str>,
1134    pub targets: HashMap<String, extension::TargetConfig>,
1135    pub env: HashMap<String, String>,
1136    pub version: Option<SharedString>,
1137    pub new_version_available_tx: Option<watch::Sender<Option<String>>>,
1138}
1139
1140impl ExternalAgentServer for LocalExtensionArchiveAgent {
1141    fn version(&self) -> Option<&SharedString> {
1142        self.version.as_ref()
1143    }
1144
1145    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1146        self.new_version_available_tx.take()
1147    }
1148
1149    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1150        self.new_version_available_tx = Some(tx);
1151    }
1152
1153    fn get_command(
1154        &self,
1155        extra_args: Vec<String>,
1156        extra_env: HashMap<String, String>,
1157        cx: &mut AsyncApp,
1158    ) -> Task<Result<AgentServerCommand>> {
1159        let fs = self.fs.clone();
1160        let http_client = self.http_client.clone();
1161        let node_runtime = self.node_runtime.clone();
1162        let project_environment = self.project_environment.downgrade();
1163        let extension_id = self.extension_id.clone();
1164        let agent_id = self.agent_id.clone();
1165        let targets = self.targets.clone();
1166        let base_env = self.env.clone();
1167        let version = self.version.clone();
1168
1169        cx.spawn(async move |cx| {
1170            // Get project environment
1171            let mut env = project_environment
1172                .update(cx, |project_environment, cx| {
1173                    project_environment.default_environment(cx)
1174                })?
1175                .await
1176                .unwrap_or_default();
1177
1178            // Merge manifest env and extra env
1179            env.extend(base_env);
1180            env.extend(extra_env);
1181
1182            let cache_key = format!("{}/{}", extension_id, agent_id);
1183            let dir = paths::external_agents_dir().join(&cache_key);
1184            fs.create_dir(&dir).await?;
1185
1186            // Determine platform key
1187            let os = if cfg!(target_os = "macos") {
1188                "darwin"
1189            } else if cfg!(target_os = "linux") {
1190                "linux"
1191            } else if cfg!(target_os = "windows") {
1192                "windows"
1193            } else {
1194                anyhow::bail!("unsupported OS");
1195            };
1196
1197            let arch = if cfg!(target_arch = "aarch64") {
1198                "aarch64"
1199            } else if cfg!(target_arch = "x86_64") {
1200                "x86_64"
1201            } else {
1202                anyhow::bail!("unsupported architecture");
1203            };
1204
1205            let platform_key = format!("{}-{}", os, arch);
1206            let target_config = targets.get(&platform_key).with_context(|| {
1207                format!(
1208                    "no target specified for platform '{}'. Available platforms: {}",
1209                    platform_key,
1210                    targets
1211                        .keys()
1212                        .map(|k| k.as_str())
1213                        .collect::<Vec<_>>()
1214                        .join(", ")
1215                )
1216            })?;
1217
1218            let archive_url = &target_config.archive;
1219            let version_dir = versioned_archive_cache_dir(
1220                &dir,
1221                version.as_ref().map(|version| version.as_ref()),
1222                archive_url,
1223            );
1224
1225            if !fs.is_dir(&version_dir).await {
1226                // Determine SHA256 for verification
1227                let sha256 = if let Some(provided_sha) = &target_config.sha256 {
1228                    // Use provided SHA256
1229                    Some(provided_sha.clone())
1230                } else if let Some(github_archive) = github_release_archive_from_url(archive_url) {
1231                    // Try to fetch SHA256 from GitHub API
1232                    if let Ok(release) = ::http_client::github::get_release_by_tag_name(
1233                        &github_archive.repo_name_with_owner,
1234                        &github_archive.tag,
1235                        http_client.clone(),
1236                    )
1237                    .await
1238                    {
1239                        // Find matching asset
1240                        if let Some(asset) = release
1241                            .assets
1242                            .iter()
1243                            .find(|a| a.name == github_archive.asset_name)
1244                        {
1245                            // Strip "sha256:" prefix if present
1246                            asset.digest.as_ref().map(|d| {
1247                                d.strip_prefix("sha256:")
1248                                    .map(|s| s.to_string())
1249                                    .unwrap_or_else(|| d.clone())
1250                            })
1251                        } else {
1252                            None
1253                        }
1254                    } else {
1255                        None
1256                    }
1257                } else {
1258                    None
1259                };
1260
1261                let asset_kind = asset_kind_for_archive_url(archive_url)?;
1262
1263                // Download and extract
1264                ::http_client::github_download::download_server_binary(
1265                    &*http_client,
1266                    archive_url,
1267                    sha256.as_deref(),
1268                    &version_dir,
1269                    asset_kind,
1270                )
1271                .await?;
1272            }
1273
1274            // Validate and resolve cmd path
1275            let cmd = &target_config.cmd;
1276
1277            let cmd_path = if cmd == "node" {
1278                // Use Zed's managed Node.js runtime
1279                node_runtime.binary_path().await?
1280            } else {
1281                if cmd.contains("..") {
1282                    anyhow::bail!("command path cannot contain '..': {}", cmd);
1283                }
1284
1285                if cmd.starts_with("./") || cmd.starts_with(".\\") {
1286                    // Relative to extraction directory
1287                    let cmd_path = version_dir.join(&cmd[2..]);
1288                    anyhow::ensure!(
1289                        fs.is_file(&cmd_path).await,
1290                        "Missing command {} after extraction",
1291                        cmd_path.to_string_lossy()
1292                    );
1293                    cmd_path
1294                } else {
1295                    // On PATH
1296                    anyhow::bail!("command must be relative (start with './'): {}", cmd);
1297                }
1298            };
1299
1300            let mut args = target_config.args.clone();
1301            args.extend(extra_args);
1302
1303            let command = AgentServerCommand {
1304                path: cmd_path,
1305                args,
1306                env: Some(env),
1307            };
1308
1309            Ok(command)
1310        })
1311    }
1312
1313    fn as_any(&self) -> &dyn Any {
1314        self
1315    }
1316
1317    fn as_any_mut(&mut self) -> &mut dyn Any {
1318        self
1319    }
1320}
1321
1322struct LocalRegistryArchiveAgent {
1323    fs: Arc<dyn Fs>,
1324    http_client: Arc<dyn HttpClient>,
1325    node_runtime: NodeRuntime,
1326    project_environment: Entity<ProjectEnvironment>,
1327    registry_id: Arc<str>,
1328    version: SharedString,
1329    targets: HashMap<String, RegistryTargetConfig>,
1330    env: HashMap<String, String>,
1331    new_version_available_tx: Option<watch::Sender<Option<String>>>,
1332}
1333
1334impl ExternalAgentServer for LocalRegistryArchiveAgent {
1335    fn version(&self) -> Option<&SharedString> {
1336        Some(&self.version)
1337    }
1338
1339    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1340        self.new_version_available_tx.take()
1341    }
1342
1343    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1344        self.new_version_available_tx = Some(tx);
1345    }
1346
1347    fn get_command(
1348        &self,
1349        extra_args: Vec<String>,
1350        extra_env: HashMap<String, String>,
1351        cx: &mut AsyncApp,
1352    ) -> Task<Result<AgentServerCommand>> {
1353        let fs = self.fs.clone();
1354        let http_client = self.http_client.clone();
1355        let node_runtime = self.node_runtime.clone();
1356        let project_environment = self.project_environment.downgrade();
1357        let registry_id = self.registry_id.clone();
1358        let targets = self.targets.clone();
1359        let settings_env = self.env.clone();
1360        let version = self.version.clone();
1361
1362        cx.spawn(async move |cx| {
1363            let mut env = project_environment
1364                .update(cx, |project_environment, cx| {
1365                    project_environment.default_environment(cx)
1366                })?
1367                .await
1368                .unwrap_or_default();
1369
1370            let dir = paths::external_agents_dir()
1371                .join("registry")
1372                .join(registry_id.as_ref());
1373            fs.create_dir(&dir).await?;
1374
1375            let os = if cfg!(target_os = "macos") {
1376                "darwin"
1377            } else if cfg!(target_os = "linux") {
1378                "linux"
1379            } else if cfg!(target_os = "windows") {
1380                "windows"
1381            } else {
1382                anyhow::bail!("unsupported OS");
1383            };
1384
1385            let arch = if cfg!(target_arch = "aarch64") {
1386                "aarch64"
1387            } else if cfg!(target_arch = "x86_64") {
1388                "x86_64"
1389            } else {
1390                anyhow::bail!("unsupported architecture");
1391            };
1392
1393            let platform_key = format!("{}-{}", os, arch);
1394            let target_config = targets.get(&platform_key).with_context(|| {
1395                format!(
1396                    "no target specified for platform '{}'. Available platforms: {}",
1397                    platform_key,
1398                    targets
1399                        .keys()
1400                        .map(|k| k.as_str())
1401                        .collect::<Vec<_>>()
1402                        .join(", ")
1403                )
1404            })?;
1405
1406            env.extend(target_config.env.clone());
1407            env.extend(extra_env);
1408            env.extend(settings_env);
1409
1410            let archive_url = &target_config.archive;
1411            let version_dir =
1412                versioned_archive_cache_dir(&dir, Some(version.as_ref()), archive_url);
1413
1414            if !fs.is_dir(&version_dir).await {
1415                let sha256 = if let Some(provided_sha) = &target_config.sha256 {
1416                    Some(provided_sha.clone())
1417                } else if let Some(github_archive) = github_release_archive_from_url(archive_url) {
1418                    if let Ok(release) = ::http_client::github::get_release_by_tag_name(
1419                        &github_archive.repo_name_with_owner,
1420                        &github_archive.tag,
1421                        http_client.clone(),
1422                    )
1423                    .await
1424                    {
1425                        if let Some(asset) = release
1426                            .assets
1427                            .iter()
1428                            .find(|a| a.name == github_archive.asset_name)
1429                        {
1430                            asset.digest.as_ref().and_then(|d| {
1431                                d.strip_prefix("sha256:")
1432                                    .map(|s| s.to_string())
1433                                    .or_else(|| Some(d.clone()))
1434                            })
1435                        } else {
1436                            None
1437                        }
1438                    } else {
1439                        None
1440                    }
1441                } else {
1442                    None
1443                };
1444
1445                let asset_kind = asset_kind_for_archive_url(archive_url)?;
1446
1447                ::http_client::github_download::download_server_binary(
1448                    &*http_client,
1449                    archive_url,
1450                    sha256.as_deref(),
1451                    &version_dir,
1452                    asset_kind,
1453                )
1454                .await?;
1455            }
1456
1457            let cmd = &target_config.cmd;
1458
1459            let cmd_path = if cmd == "node" {
1460                node_runtime.binary_path().await?
1461            } else {
1462                if cmd.contains("..") {
1463                    anyhow::bail!("command path cannot contain '..': {}", cmd);
1464                }
1465
1466                if cmd.starts_with("./") || cmd.starts_with(".\\") {
1467                    let cmd_path = version_dir.join(&cmd[2..]);
1468                    anyhow::ensure!(
1469                        fs.is_file(&cmd_path).await,
1470                        "Missing command {} after extraction",
1471                        cmd_path.to_string_lossy()
1472                    );
1473                    cmd_path
1474                } else {
1475                    anyhow::bail!("command must be relative (start with './'): {}", cmd);
1476                }
1477            };
1478
1479            let mut args = target_config.args.clone();
1480            args.extend(extra_args);
1481
1482            let command = AgentServerCommand {
1483                path: cmd_path,
1484                args,
1485                env: Some(env),
1486            };
1487
1488            Ok(command)
1489        })
1490    }
1491
1492    fn as_any(&self) -> &dyn Any {
1493        self
1494    }
1495
1496    fn as_any_mut(&mut self) -> &mut dyn Any {
1497        self
1498    }
1499}
1500
1501struct LocalRegistryNpxAgent {
1502    fs: Arc<dyn Fs>,
1503    node_runtime: NodeRuntime,
1504    project_environment: Entity<ProjectEnvironment>,
1505    version: SharedString,
1506    package: SharedString,
1507    args: Vec<String>,
1508    distribution_env: HashMap<String, String>,
1509    settings_env: HashMap<String, String>,
1510    new_version_available_tx: Option<watch::Sender<Option<String>>>,
1511}
1512
1513impl ExternalAgentServer for LocalRegistryNpxAgent {
1514    fn version(&self) -> Option<&SharedString> {
1515        Some(&self.version)
1516    }
1517
1518    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1519        self.new_version_available_tx.take()
1520    }
1521
1522    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1523        self.new_version_available_tx = Some(tx);
1524    }
1525
1526    fn get_command(
1527        &self,
1528        extra_args: Vec<String>,
1529        extra_env: HashMap<String, String>,
1530        cx: &mut AsyncApp,
1531    ) -> Task<Result<AgentServerCommand>> {
1532        let fs = self.fs.clone();
1533        let node_runtime = self.node_runtime.clone();
1534        let project_environment = self.project_environment.downgrade();
1535        let package = self.package.clone();
1536        let args = self.args.clone();
1537        let distribution_env = self.distribution_env.clone();
1538        let settings_env = self.settings_env.clone();
1539
1540        cx.spawn(async move |cx| {
1541            let mut env = project_environment
1542                .update(cx, |project_environment, cx| {
1543                    project_environment.default_environment(cx)
1544                })?
1545                .await
1546                .unwrap_or_default();
1547
1548            let prefix_dir = paths::external_agents_dir()
1549                .join("registry")
1550                .join("npx")
1551                .join(Path::new(package.as_str()));
1552            fs.create_dir(&prefix_dir).await?;
1553
1554            let mut exec_args = vec!["--yes".to_string(), "--".to_string(), package.to_string()];
1555            exec_args.extend(args);
1556
1557            let npm_command = node_runtime
1558                .npm_command(
1559                    Some(&prefix_dir),
1560                    "exec",
1561                    &exec_args.iter().map(|a| a.as_str()).collect::<Vec<_>>(),
1562                )
1563                .await?;
1564
1565            env.extend(npm_command.env);
1566            env.extend(distribution_env);
1567            env.extend(extra_env);
1568            env.extend(settings_env);
1569
1570            let mut args = npm_command.args;
1571            args.extend(extra_args);
1572
1573            let command = AgentServerCommand {
1574                path: npm_command.path,
1575                args,
1576                env: Some(env),
1577            };
1578
1579            Ok(command)
1580        })
1581    }
1582
1583    fn as_any(&self) -> &dyn Any {
1584        self
1585    }
1586
1587    fn as_any_mut(&mut self) -> &mut dyn Any {
1588        self
1589    }
1590}
1591
1592struct LocalCustomAgent {
1593    project_environment: Entity<ProjectEnvironment>,
1594    command: AgentServerCommand,
1595}
1596
1597impl ExternalAgentServer for LocalCustomAgent {
1598    fn get_command(
1599        &self,
1600        extra_args: Vec<String>,
1601        extra_env: HashMap<String, String>,
1602        cx: &mut AsyncApp,
1603    ) -> Task<Result<AgentServerCommand>> {
1604        let mut command = self.command.clone();
1605        let project_environment = self.project_environment.downgrade();
1606        cx.spawn(async move |cx| {
1607            let mut env = project_environment
1608                .update(cx, |project_environment, cx| {
1609                    project_environment.default_environment(cx)
1610                })?
1611                .await
1612                .unwrap_or_default();
1613            env.extend(command.env.unwrap_or_default());
1614            env.extend(extra_env);
1615            command.env = Some(env);
1616            command.args.extend(extra_args);
1617            Ok(command)
1618        })
1619    }
1620
1621    fn as_any(&self) -> &dyn Any {
1622        self
1623    }
1624
1625    fn as_any_mut(&mut self) -> &mut dyn Any {
1626        self
1627    }
1628}
1629
1630#[derive(Default, Clone, JsonSchema, Debug, PartialEq, RegisterSetting)]
1631pub struct AllAgentServersSettings(pub HashMap<String, CustomAgentServerSettings>);
1632
1633impl std::ops::Deref for AllAgentServersSettings {
1634    type Target = HashMap<String, CustomAgentServerSettings>;
1635
1636    fn deref(&self) -> &Self::Target {
1637        &self.0
1638    }
1639}
1640
1641impl std::ops::DerefMut for AllAgentServersSettings {
1642    fn deref_mut(&mut self) -> &mut Self::Target {
1643        &mut self.0
1644    }
1645}
1646
1647impl AllAgentServersSettings {
1648    pub fn has_registry_agents(&self) -> bool {
1649        self.values()
1650            .any(|s| matches!(s, CustomAgentServerSettings::Registry { .. }))
1651    }
1652}
1653
1654#[derive(Clone, JsonSchema, Debug, PartialEq)]
1655pub enum CustomAgentServerSettings {
1656    Custom {
1657        command: AgentServerCommand,
1658        /// The default mode to use for this agent.
1659        ///
1660        /// Note: Not only all agents support modes.
1661        ///
1662        /// Default: None
1663        default_mode: Option<String>,
1664        /// The default model to use for this agent.
1665        ///
1666        /// This should be the model ID as reported by the agent.
1667        ///
1668        /// Default: None
1669        default_model: Option<String>,
1670        /// The favorite models for this agent.
1671        ///
1672        /// Default: []
1673        favorite_models: Vec<String>,
1674        /// Default values for session config options.
1675        ///
1676        /// This is a map from config option ID to value ID.
1677        ///
1678        /// Default: {}
1679        default_config_options: HashMap<String, String>,
1680        /// Favorited values for session config options.
1681        ///
1682        /// This is a map from config option ID to a list of favorited value IDs.
1683        ///
1684        /// Default: {}
1685        favorite_config_option_values: HashMap<String, Vec<String>>,
1686    },
1687    Extension {
1688        /// Additional environment variables to pass to the agent.
1689        ///
1690        /// Default: {}
1691        env: HashMap<String, String>,
1692        /// The default mode to use for this agent.
1693        ///
1694        /// Note: Not only all agents support modes.
1695        ///
1696        /// Default: None
1697        default_mode: Option<String>,
1698        /// The default model to use for this agent.
1699        ///
1700        /// This should be the model ID as reported by the agent.
1701        ///
1702        /// Default: None
1703        default_model: Option<String>,
1704        /// The favorite models for this agent.
1705        ///
1706        /// Default: []
1707        favorite_models: Vec<String>,
1708        /// Default values for session config options.
1709        ///
1710        /// This is a map from config option ID to value ID.
1711        ///
1712        /// Default: {}
1713        default_config_options: HashMap<String, String>,
1714        /// Favorited values for session config options.
1715        ///
1716        /// This is a map from config option ID to a list of favorited value IDs.
1717        ///
1718        /// Default: {}
1719        favorite_config_option_values: HashMap<String, Vec<String>>,
1720    },
1721    Registry {
1722        /// Additional environment variables to pass to the agent.
1723        ///
1724        /// Default: {}
1725        env: HashMap<String, String>,
1726        /// The default mode to use for this agent.
1727        ///
1728        /// Note: Not only all agents support modes.
1729        ///
1730        /// Default: None
1731        default_mode: Option<String>,
1732        /// The default model to use for this agent.
1733        ///
1734        /// This should be the model ID as reported by the agent.
1735        ///
1736        /// Default: None
1737        default_model: Option<String>,
1738        /// The favorite models for this agent.
1739        ///
1740        /// Default: []
1741        favorite_models: Vec<String>,
1742        /// Default values for session config options.
1743        ///
1744        /// This is a map from config option ID to value ID.
1745        ///
1746        /// Default: {}
1747        default_config_options: HashMap<String, String>,
1748        /// Favorited values for session config options.
1749        ///
1750        /// This is a map from config option ID to a list of favorited value IDs.
1751        ///
1752        /// Default: {}
1753        favorite_config_option_values: HashMap<String, Vec<String>>,
1754    },
1755}
1756
1757impl CustomAgentServerSettings {
1758    pub fn command(&self) -> Option<&AgentServerCommand> {
1759        match self {
1760            CustomAgentServerSettings::Custom { command, .. } => Some(command),
1761            CustomAgentServerSettings::Extension { .. }
1762            | CustomAgentServerSettings::Registry { .. } => None,
1763        }
1764    }
1765
1766    pub fn default_mode(&self) -> Option<&str> {
1767        match self {
1768            CustomAgentServerSettings::Custom { default_mode, .. }
1769            | CustomAgentServerSettings::Extension { default_mode, .. }
1770            | CustomAgentServerSettings::Registry { default_mode, .. } => default_mode.as_deref(),
1771        }
1772    }
1773
1774    pub fn default_model(&self) -> Option<&str> {
1775        match self {
1776            CustomAgentServerSettings::Custom { default_model, .. }
1777            | CustomAgentServerSettings::Extension { default_model, .. }
1778            | CustomAgentServerSettings::Registry { default_model, .. } => default_model.as_deref(),
1779        }
1780    }
1781
1782    pub fn favorite_models(&self) -> &[String] {
1783        match self {
1784            CustomAgentServerSettings::Custom {
1785                favorite_models, ..
1786            }
1787            | CustomAgentServerSettings::Extension {
1788                favorite_models, ..
1789            }
1790            | CustomAgentServerSettings::Registry {
1791                favorite_models, ..
1792            } => favorite_models,
1793        }
1794    }
1795
1796    pub fn default_config_option(&self, config_id: &str) -> Option<&str> {
1797        match self {
1798            CustomAgentServerSettings::Custom {
1799                default_config_options,
1800                ..
1801            }
1802            | CustomAgentServerSettings::Extension {
1803                default_config_options,
1804                ..
1805            }
1806            | CustomAgentServerSettings::Registry {
1807                default_config_options,
1808                ..
1809            } => default_config_options.get(config_id).map(|s| s.as_str()),
1810        }
1811    }
1812
1813    pub fn favorite_config_option_values(&self, config_id: &str) -> Option<&[String]> {
1814        match self {
1815            CustomAgentServerSettings::Custom {
1816                favorite_config_option_values,
1817                ..
1818            }
1819            | CustomAgentServerSettings::Extension {
1820                favorite_config_option_values,
1821                ..
1822            }
1823            | CustomAgentServerSettings::Registry {
1824                favorite_config_option_values,
1825                ..
1826            } => favorite_config_option_values
1827                .get(config_id)
1828                .map(|v| v.as_slice()),
1829        }
1830    }
1831}
1832
1833impl From<settings::CustomAgentServerSettings> for CustomAgentServerSettings {
1834    fn from(value: settings::CustomAgentServerSettings) -> Self {
1835        match value {
1836            settings::CustomAgentServerSettings::Custom {
1837                path,
1838                args,
1839                env,
1840                default_mode,
1841                default_model,
1842                favorite_models,
1843                default_config_options,
1844                favorite_config_option_values,
1845            } => CustomAgentServerSettings::Custom {
1846                command: AgentServerCommand {
1847                    path: PathBuf::from(shellexpand::tilde(&path.to_string_lossy()).as_ref()),
1848                    args,
1849                    env: Some(env),
1850                },
1851                default_mode,
1852                default_model,
1853                favorite_models,
1854                default_config_options,
1855                favorite_config_option_values,
1856            },
1857            settings::CustomAgentServerSettings::Extension {
1858                env,
1859                default_mode,
1860                default_model,
1861                default_config_options,
1862                favorite_models,
1863                favorite_config_option_values,
1864            } => CustomAgentServerSettings::Extension {
1865                env,
1866                default_mode,
1867                default_model,
1868                default_config_options,
1869                favorite_models,
1870                favorite_config_option_values,
1871            },
1872            settings::CustomAgentServerSettings::Registry {
1873                env,
1874                default_mode,
1875                default_model,
1876                default_config_options,
1877                favorite_models,
1878                favorite_config_option_values,
1879            } => CustomAgentServerSettings::Registry {
1880                env,
1881                default_mode,
1882                default_model,
1883                default_config_options,
1884                favorite_models,
1885                favorite_config_option_values,
1886            },
1887        }
1888    }
1889}
1890
1891impl settings::Settings for AllAgentServersSettings {
1892    fn from_settings(content: &settings::SettingsContent) -> Self {
1893        let agent_settings = content.agent_servers.clone().unwrap();
1894        Self(
1895            agent_settings
1896                .0
1897                .into_iter()
1898                .map(|(k, v)| (k, v.into()))
1899                .collect(),
1900        )
1901    }
1902}
1903
1904#[cfg(test)]
1905mod tests {
1906    use super::*;
1907    use crate::agent_registry_store::{
1908        AgentRegistryStore, RegistryAgent, RegistryAgentMetadata, RegistryNpxAgent,
1909    };
1910    use crate::worktree_store::{WorktreeIdCounter, WorktreeStore};
1911    use gpui::{AppContext as _, TestAppContext};
1912    use node_runtime::NodeRuntime;
1913    use settings::Settings as _;
1914
1915    fn make_npx_agent(id: &str, version: &str) -> RegistryAgent {
1916        let id = SharedString::from(id.to_string());
1917        RegistryAgent::Npx(RegistryNpxAgent {
1918            metadata: RegistryAgentMetadata {
1919                id: AgentId::new(id.clone()),
1920                name: id.clone(),
1921                description: SharedString::from(""),
1922                version: SharedString::from(version.to_string()),
1923                repository: None,
1924                website: None,
1925                icon_path: None,
1926            },
1927            package: id,
1928            args: Vec::new(),
1929            env: HashMap::default(),
1930        })
1931    }
1932
1933    fn init_test_settings(cx: &mut TestAppContext) {
1934        cx.update(|cx| {
1935            let settings_store = SettingsStore::test(cx);
1936            cx.set_global(settings_store);
1937        });
1938    }
1939
1940    fn init_registry(
1941        cx: &mut TestAppContext,
1942        agents: Vec<RegistryAgent>,
1943    ) -> gpui::Entity<AgentRegistryStore> {
1944        cx.update(|cx| AgentRegistryStore::init_test_global(cx, agents))
1945    }
1946
1947    fn set_registry_settings(cx: &mut TestAppContext, agent_names: &[&str]) {
1948        cx.update(|cx| {
1949            AllAgentServersSettings::override_global(
1950                AllAgentServersSettings(
1951                    agent_names
1952                        .iter()
1953                        .map(|name| {
1954                            (
1955                                name.to_string(),
1956                                settings::CustomAgentServerSettings::Registry {
1957                                    env: HashMap::default(),
1958                                    default_mode: None,
1959                                    default_model: None,
1960                                    favorite_models: Vec::new(),
1961                                    default_config_options: HashMap::default(),
1962                                    favorite_config_option_values: HashMap::default(),
1963                                }
1964                                .into(),
1965                            )
1966                        })
1967                        .collect(),
1968                ),
1969                cx,
1970            );
1971        });
1972    }
1973
1974    fn create_agent_server_store(cx: &mut TestAppContext) -> gpui::Entity<AgentServerStore> {
1975        cx.update(|cx| {
1976            let fs: Arc<dyn Fs> = fs::FakeFs::new(cx.background_executor().clone());
1977            let worktree_store =
1978                cx.new(|cx| WorktreeStore::local(false, fs.clone(), WorktreeIdCounter::get(cx)));
1979            let project_environment = cx.new(|cx| {
1980                crate::ProjectEnvironment::new(None, worktree_store.downgrade(), None, false, cx)
1981            });
1982            let http_client = http_client::FakeHttpClient::with_404_response();
1983
1984            cx.new(|cx| {
1985                AgentServerStore::local(
1986                    NodeRuntime::unavailable(),
1987                    fs,
1988                    project_environment,
1989                    http_client,
1990                    cx,
1991                )
1992            })
1993        })
1994    }
1995
1996    #[test]
1997    fn detects_supported_archive_suffixes() {
1998        assert!(matches!(
1999            asset_kind_for_archive_url("https://example.com/agent.zip"),
2000            Ok(AssetKind::Zip)
2001        ));
2002        assert!(matches!(
2003            asset_kind_for_archive_url("https://example.com/agent.zip?download=1"),
2004            Ok(AssetKind::Zip)
2005        ));
2006        assert!(matches!(
2007            asset_kind_for_archive_url("https://example.com/agent.tar.gz"),
2008            Ok(AssetKind::TarGz)
2009        ));
2010        assert!(matches!(
2011            asset_kind_for_archive_url("https://example.com/agent.tar.gz?download=1#latest"),
2012            Ok(AssetKind::TarGz)
2013        ));
2014        assert!(matches!(
2015            asset_kind_for_archive_url("https://example.com/agent.tgz"),
2016            Ok(AssetKind::TarGz)
2017        ));
2018        assert!(matches!(
2019            asset_kind_for_archive_url("https://example.com/agent.tgz#download"),
2020            Ok(AssetKind::TarGz)
2021        ));
2022        assert!(matches!(
2023            asset_kind_for_archive_url("https://example.com/agent.tar.bz2"),
2024            Ok(AssetKind::TarBz2)
2025        ));
2026        assert!(matches!(
2027            asset_kind_for_archive_url("https://example.com/agent.tar.bz2?download=1"),
2028            Ok(AssetKind::TarBz2)
2029        ));
2030        assert!(matches!(
2031            asset_kind_for_archive_url("https://example.com/agent.tbz2"),
2032            Ok(AssetKind::TarBz2)
2033        ));
2034        assert!(matches!(
2035            asset_kind_for_archive_url("https://example.com/agent.tbz2#download"),
2036            Ok(AssetKind::TarBz2)
2037        ));
2038    }
2039
2040    #[test]
2041    fn parses_github_release_archive_urls() {
2042        let github_archive = github_release_archive_from_url(
2043            "https://github.com/owner/repo/releases/download/release%2F2.3.5/agent.tar.bz2?download=1",
2044        )
2045        .unwrap();
2046
2047        assert_eq!(github_archive.repo_name_with_owner, "owner/repo");
2048        assert_eq!(github_archive.tag, "release/2.3.5");
2049        assert_eq!(github_archive.asset_name, "agent.tar.bz2");
2050    }
2051
2052    #[test]
2053    fn rejects_unsupported_archive_suffixes() {
2054        let error = asset_kind_for_archive_url("https://example.com/agent.tar.xz")
2055            .err()
2056            .map(|error| error.to_string());
2057
2058        assert_eq!(
2059            error,
2060            Some("unsupported archive type in URL: https://example.com/agent.tar.xz".to_string()),
2061        );
2062    }
2063
2064    #[test]
2065    fn versioned_archive_cache_dir_includes_version_before_url_hash() {
2066        let slash_version_dir = versioned_archive_cache_dir(
2067            Path::new("/tmp/agents"),
2068            Some("release/2.3.5"),
2069            "https://example.com/agent.zip",
2070        );
2071        let colon_version_dir = versioned_archive_cache_dir(
2072            Path::new("/tmp/agents"),
2073            Some("release:2.3.5"),
2074            "https://example.com/agent.zip",
2075        );
2076        let file_name = slash_version_dir
2077            .file_name()
2078            .and_then(|name| name.to_str())
2079            .expect("cache directory should have a file name");
2080
2081        assert!(file_name.starts_with("v_release-2.3.5_"));
2082        assert_ne!(slash_version_dir, colon_version_dir);
2083    }
2084
2085    #[gpui::test]
2086    fn test_version_change_sends_notification(cx: &mut TestAppContext) {
2087        init_test_settings(cx);
2088        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2089        set_registry_settings(cx, &["test-agent"]);
2090        let store = create_agent_server_store(cx);
2091
2092        // Verify the agent was registered with version 1.0.0.
2093        store.read_with(cx, |store, _| {
2094            let entry = store
2095                .external_agents
2096                .get(&AgentId::new("test-agent"))
2097                .expect("agent should be registered");
2098            assert_eq!(
2099                entry.server.version().map(|v| v.to_string()),
2100                Some("1.0.0".to_string())
2101            );
2102        });
2103
2104        // Set up a watch channel and store the tx on the agent.
2105        let (tx, mut rx) = watch::channel::<Option<String>>(None);
2106        store.update(cx, |store, _| {
2107            let entry = store
2108                .external_agents
2109                .get_mut(&AgentId::new("test-agent"))
2110                .expect("agent should be registered");
2111            entry.server.set_new_version_available_tx(tx);
2112        });
2113
2114        // Update the registry to version 2.0.0.
2115        registry.update(cx, |store, cx| {
2116            store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx);
2117        });
2118        cx.run_until_parked();
2119
2120        // The watch channel should have received the new version.
2121        assert_eq!(rx.borrow().as_deref(), Some("2.0.0"));
2122    }
2123
2124    #[gpui::test]
2125    fn test_same_version_preserves_tx(cx: &mut TestAppContext) {
2126        init_test_settings(cx);
2127        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2128        set_registry_settings(cx, &["test-agent"]);
2129        let store = create_agent_server_store(cx);
2130
2131        let (tx, mut rx) = watch::channel::<Option<String>>(None);
2132        store.update(cx, |store, _| {
2133            let entry = store
2134                .external_agents
2135                .get_mut(&AgentId::new("test-agent"))
2136                .expect("agent should be registered");
2137            entry.server.set_new_version_available_tx(tx);
2138        });
2139
2140        // "Refresh" the registry with the same version.
2141        registry.update(cx, |store, cx| {
2142            store.set_agents(vec![make_npx_agent("test-agent", "1.0.0")], cx);
2143        });
2144        cx.run_until_parked();
2145
2146        // No notification should have been sent.
2147        assert_eq!(rx.borrow().as_deref(), None);
2148
2149        // The tx should have been transferred to the rebuilt agent entry.
2150        store.update(cx, |store, _| {
2151            let entry = store
2152                .external_agents
2153                .get_mut(&AgentId::new("test-agent"))
2154                .expect("agent should be registered");
2155            assert!(
2156                entry.server.take_new_version_available_tx().is_some(),
2157                "tx should have been transferred to the rebuilt agent"
2158            );
2159        });
2160    }
2161
2162    #[gpui::test]
2163    fn test_no_tx_stored_does_not_panic_on_version_change(cx: &mut TestAppContext) {
2164        init_test_settings(cx);
2165        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2166        set_registry_settings(cx, &["test-agent"]);
2167        let _store = create_agent_server_store(cx);
2168
2169        // Update the registry without having stored any tx — should not panic.
2170        registry.update(cx, |store, cx| {
2171            store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx);
2172        });
2173        cx.run_until_parked();
2174    }
2175
2176    #[gpui::test]
2177    fn test_multiple_agents_independent_notifications(cx: &mut TestAppContext) {
2178        init_test_settings(cx);
2179        let registry = init_registry(
2180            cx,
2181            vec![
2182                make_npx_agent("agent-a", "1.0.0"),
2183                make_npx_agent("agent-b", "3.0.0"),
2184            ],
2185        );
2186        set_registry_settings(cx, &["agent-a", "agent-b"]);
2187        let store = create_agent_server_store(cx);
2188
2189        let (tx_a, mut rx_a) = watch::channel::<Option<String>>(None);
2190        let (tx_b, mut rx_b) = watch::channel::<Option<String>>(None);
2191        store.update(cx, |store, _| {
2192            store
2193                .external_agents
2194                .get_mut(&AgentId::new("agent-a"))
2195                .expect("agent-a should be registered")
2196                .server
2197                .set_new_version_available_tx(tx_a);
2198            store
2199                .external_agents
2200                .get_mut(&AgentId::new("agent-b"))
2201                .expect("agent-b should be registered")
2202                .server
2203                .set_new_version_available_tx(tx_b);
2204        });
2205
2206        // Update only agent-a to a new version; agent-b stays the same.
2207        registry.update(cx, |store, cx| {
2208            store.set_agents(
2209                vec![
2210                    make_npx_agent("agent-a", "2.0.0"),
2211                    make_npx_agent("agent-b", "3.0.0"),
2212                ],
2213                cx,
2214            );
2215        });
2216        cx.run_until_parked();
2217
2218        // agent-a should have received a notification.
2219        assert_eq!(rx_a.borrow().as_deref(), Some("2.0.0"));
2220
2221        // agent-b should NOT have received a notification.
2222        assert_eq!(rx_b.borrow().as_deref(), None);
2223
2224        // agent-b's tx should have been transferred.
2225        store.update(cx, |store, _| {
2226            assert!(
2227                store
2228                    .external_agents
2229                    .get_mut(&AgentId::new("agent-b"))
2230                    .expect("agent-b should be registered")
2231                    .server
2232                    .take_new_version_available_tx()
2233                    .is_some(),
2234                "agent-b tx should have been transferred"
2235            );
2236        });
2237    }
2238}