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                                        node_runtime: node_runtime.clone(),
 572                                        project_environment: project_environment.clone(),
 573                                        version: agent.metadata.version.clone(),
 574                                        package: agent.package.clone(),
 575                                        args: agent.args.clone(),
 576                                        distribution_env: agent.env.clone(),
 577                                        settings_env: env.clone(),
 578                                        new_version_available_tx: None,
 579                                    })
 580                                        as Box<dyn ExternalAgentServer>,
 581                                    ExternalAgentSource::Registry,
 582                                    agent.metadata.icon_path.clone(),
 583                                    Some(agent.metadata.name.clone()),
 584                                ),
 585                            );
 586                        }
 587                    }
 588                }
 589                CustomAgentServerSettings::Extension { .. } => {}
 590            }
 591        }
 592
 593        // For each rebuilt versioned agent, compare the version. If it
 594        // changed, notify the active connection to reconnect. Otherwise,
 595        // transfer the channel to the new entry so future updates can use it.
 596        for (name, entry) in &mut self.external_agents {
 597            let Some((old_version, mut tx)) = old_versioned_agents.remove(name) else {
 598                continue;
 599            };
 600            let Some(new_version) = entry.server.version() else {
 601                continue;
 602            };
 603
 604            if new_version != &old_version {
 605                tx.send(Some(new_version.to_string())).ok();
 606            } else {
 607                entry.server.set_new_version_available_tx(tx);
 608            }
 609        }
 610
 611        *old_settings = Some(new_settings);
 612
 613        if let Some((project_id, downstream_client)) = downstream_client {
 614            downstream_client
 615                .send(proto::ExternalAgentsUpdated {
 616                    project_id: *project_id,
 617                    names: self
 618                        .external_agents
 619                        .keys()
 620                        .map(|name| name.to_string())
 621                        .collect(),
 622                })
 623                .log_err();
 624        }
 625        cx.emit(AgentServersUpdated);
 626    }
 627
 628    pub fn node_runtime(&self) -> Option<NodeRuntime> {
 629        match &self.state {
 630            AgentServerStoreState::Local { node_runtime, .. } => Some(node_runtime.clone()),
 631            _ => None,
 632        }
 633    }
 634
 635    pub fn local(
 636        node_runtime: NodeRuntime,
 637        fs: Arc<dyn Fs>,
 638        project_environment: Entity<ProjectEnvironment>,
 639        http_client: Arc<dyn HttpClient>,
 640        cx: &mut Context<Self>,
 641    ) -> Self {
 642        let mut subscriptions = vec![cx.observe_global::<SettingsStore>(|this, cx| {
 643            this.agent_servers_settings_changed(cx);
 644        })];
 645        if let Some(registry_store) = AgentRegistryStore::try_global(cx) {
 646            subscriptions.push(cx.observe(&registry_store, |this, _, cx| {
 647                this.reregister_agents(cx);
 648            }));
 649        }
 650        let mut this = Self {
 651            state: AgentServerStoreState::Local {
 652                node_runtime,
 653                fs,
 654                project_environment,
 655                http_client,
 656                downstream_client: None,
 657                settings: None,
 658                extension_agents: vec![],
 659                _subscriptions: subscriptions,
 660            },
 661            external_agents: HashMap::default(),
 662        };
 663        if let Some(_events) = extension::ExtensionEvents::try_global(cx) {}
 664        this.agent_servers_settings_changed(cx);
 665        this
 666    }
 667
 668    pub(crate) fn remote(
 669        project_id: u64,
 670        upstream_client: Entity<RemoteClient>,
 671        worktree_store: Entity<WorktreeStore>,
 672    ) -> Self {
 673        Self {
 674            state: AgentServerStoreState::Remote {
 675                project_id,
 676                upstream_client,
 677                worktree_store,
 678            },
 679            external_agents: HashMap::default(),
 680        }
 681    }
 682
 683    pub fn collab() -> Self {
 684        Self {
 685            state: AgentServerStoreState::Collab,
 686            external_agents: HashMap::default(),
 687        }
 688    }
 689
 690    pub fn shared(&mut self, project_id: u64, client: AnyProtoClient, cx: &mut Context<Self>) {
 691        match &mut self.state {
 692            AgentServerStoreState::Local {
 693                downstream_client, ..
 694            } => {
 695                *downstream_client = Some((project_id, client.clone()));
 696                // Send the current list of external agents downstream, but only after a delay,
 697                // to avoid having the message arrive before the downstream project's agent server store
 698                // sets up its handlers.
 699                cx.spawn(async move |this, cx| {
 700                    cx.background_executor().timer(Duration::from_secs(1)).await;
 701                    let names = this.update(cx, |this, _| {
 702                        this.external_agents()
 703                            .map(|name| name.to_string())
 704                            .collect()
 705                    })?;
 706                    client
 707                        .send(proto::ExternalAgentsUpdated { project_id, names })
 708                        .log_err();
 709                    anyhow::Ok(())
 710                })
 711                .detach();
 712            }
 713            AgentServerStoreState::Remote { .. } => {
 714                debug_panic!(
 715                    "external agents over collab not implemented, remote project should not be shared"
 716                );
 717            }
 718            AgentServerStoreState::Collab => {
 719                debug_panic!("external agents over collab not implemented, should not be shared");
 720            }
 721        }
 722    }
 723
 724    pub fn get_external_agent(
 725        &mut self,
 726        name: &AgentId,
 727    ) -> Option<&mut (dyn ExternalAgentServer + 'static)> {
 728        self.external_agents
 729            .get_mut(name)
 730            .map(|entry| entry.server.as_mut())
 731    }
 732
 733    pub fn no_browser(&self) -> bool {
 734        match &self.state {
 735            AgentServerStoreState::Local {
 736                downstream_client, ..
 737            } => downstream_client
 738                .as_ref()
 739                .is_some_and(|(_, client)| !client.has_wsl_interop()),
 740            _ => false,
 741        }
 742    }
 743
 744    pub fn has_external_agents(&self) -> bool {
 745        !self.external_agents.is_empty()
 746    }
 747
 748    pub fn external_agents(&self) -> impl Iterator<Item = &AgentId> {
 749        self.external_agents.keys()
 750    }
 751
 752    async fn handle_get_agent_server_command(
 753        this: Entity<Self>,
 754        envelope: TypedEnvelope<proto::GetAgentServerCommand>,
 755        mut cx: AsyncApp,
 756    ) -> Result<proto::AgentServerCommand> {
 757        let command = this
 758            .update(&mut cx, |this, cx| {
 759                let AgentServerStoreState::Local {
 760                    downstream_client, ..
 761                } = &this.state
 762                else {
 763                    debug_panic!("should not receive GetAgentServerCommand in a non-local project");
 764                    bail!("unexpected GetAgentServerCommand request in a non-local project");
 765                };
 766                let no_browser = this.no_browser();
 767                let agent = this
 768                    .external_agents
 769                    .get_mut(&*envelope.payload.name)
 770                    .map(|entry| entry.server.as_mut())
 771                    .with_context(|| format!("agent `{}` not found", envelope.payload.name))?;
 772                let new_version_available_tx =
 773                    downstream_client
 774                        .clone()
 775                        .map(|(project_id, downstream_client)| {
 776                            let (new_version_available_tx, mut new_version_available_rx) =
 777                                watch::channel(None);
 778                            cx.spawn({
 779                                let name = envelope.payload.name.clone();
 780                                async move |_, _| {
 781                                    if let Some(version) =
 782                                        new_version_available_rx.recv().await.ok().flatten()
 783                                    {
 784                                        downstream_client.send(
 785                                            proto::NewExternalAgentVersionAvailable {
 786                                                project_id,
 787                                                name: name.clone(),
 788                                                version,
 789                                            },
 790                                        )?;
 791                                    }
 792                                    anyhow::Ok(())
 793                                }
 794                            })
 795                            .detach_and_log_err(cx);
 796                            new_version_available_tx
 797                        });
 798                let mut extra_env = HashMap::default();
 799                if no_browser {
 800                    extra_env.insert("NO_BROWSER".to_owned(), "1".to_owned());
 801                }
 802                if let Some(new_version_available_tx) = new_version_available_tx {
 803                    agent.set_new_version_available_tx(new_version_available_tx);
 804                }
 805                anyhow::Ok(agent.get_command(vec![], extra_env, &mut cx.to_async()))
 806            })?
 807            .await?;
 808        Ok(proto::AgentServerCommand {
 809            path: command.path.to_string_lossy().into_owned(),
 810            args: command.args,
 811            env: command
 812                .env
 813                .map(|env| env.into_iter().collect())
 814                .unwrap_or_default(),
 815            root_dir: envelope
 816                .payload
 817                .root_dir
 818                .unwrap_or_else(|| paths::home_dir().to_string_lossy().to_string()),
 819            login: None,
 820        })
 821    }
 822
 823    async fn handle_external_agents_updated(
 824        this: Entity<Self>,
 825        envelope: TypedEnvelope<proto::ExternalAgentsUpdated>,
 826        mut cx: AsyncApp,
 827    ) -> Result<()> {
 828        this.update(&mut cx, |this, cx| {
 829            let AgentServerStoreState::Remote {
 830                project_id,
 831                upstream_client,
 832                worktree_store,
 833            } = &this.state
 834            else {
 835                debug_panic!(
 836                    "handle_external_agents_updated should not be called for a non-remote project"
 837                );
 838                bail!("unexpected ExternalAgentsUpdated message")
 839            };
 840
 841            let mut previous_entries = std::mem::take(&mut this.external_agents);
 842            let mut new_version_available_txs = HashMap::default();
 843            let mut metadata = HashMap::default();
 844
 845            for (name, mut entry) in previous_entries.drain() {
 846                if let Some(tx) = entry.server.take_new_version_available_tx() {
 847                    new_version_available_txs.insert(name.clone(), tx);
 848                }
 849
 850                metadata.insert(name, (entry.icon, entry.display_name, entry.source));
 851            }
 852
 853            this.external_agents = envelope
 854                .payload
 855                .names
 856                .into_iter()
 857                .map(|name| {
 858                    let agent_id = AgentId(name.into());
 859                    let (icon, display_name, source) = metadata
 860                        .remove(&agent_id)
 861                        .or_else(|| {
 862                            AgentRegistryStore::try_global(cx)
 863                                .and_then(|store| store.read(cx).agent(&agent_id))
 864                                .map(|s| {
 865                                    (
 866                                        s.icon_path().cloned(),
 867                                        Some(s.name().clone()),
 868                                        ExternalAgentSource::Registry,
 869                                    )
 870                                })
 871                        })
 872                        .unwrap_or((None, None, ExternalAgentSource::default()));
 873                    let agent = RemoteExternalAgentServer {
 874                        project_id: *project_id,
 875                        upstream_client: upstream_client.clone(),
 876                        worktree_store: worktree_store.clone(),
 877                        name: agent_id.clone(),
 878                        new_version_available_tx: new_version_available_txs.remove(&agent_id),
 879                    };
 880                    (
 881                        agent_id,
 882                        ExternalAgentEntry::new(
 883                            Box::new(agent) as Box<dyn ExternalAgentServer>,
 884                            source,
 885                            icon,
 886                            display_name,
 887                        ),
 888                    )
 889                })
 890                .collect();
 891            cx.emit(AgentServersUpdated);
 892            Ok(())
 893        })
 894    }
 895
 896    async fn handle_external_extension_agents_updated(
 897        this: Entity<Self>,
 898        envelope: TypedEnvelope<proto::ExternalExtensionAgentsUpdated>,
 899        mut cx: AsyncApp,
 900    ) -> Result<()> {
 901        this.update(&mut cx, |this, cx| {
 902            let AgentServerStoreState::Local {
 903                extension_agents, ..
 904            } = &mut this.state
 905            else {
 906                panic!(
 907                    "handle_external_extension_agents_updated \
 908                    should not be called for a non-remote project"
 909                );
 910            };
 911
 912            extension_agents.clear();
 913            for ExternalExtensionAgent {
 914                name,
 915                icon_path,
 916                extension_id,
 917                targets,
 918                env,
 919                version,
 920            } in envelope.payload.agents
 921            {
 922                extension_agents.push(ExtensionAgentEntry {
 923                    agent_name: Arc::from(&*name),
 924                    extension_id,
 925                    targets: targets
 926                        .into_iter()
 927                        .map(|(k, v)| (k, extension::TargetConfig::from_proto(v)))
 928                        .collect(),
 929                    env: env.into_iter().collect(),
 930                    icon_path,
 931                    display_name: None,
 932                    version: version.map(SharedString::from),
 933                });
 934            }
 935
 936            this.reregister_agents(cx);
 937            cx.emit(AgentServersUpdated);
 938            Ok(())
 939        })
 940    }
 941
 942    async fn handle_new_version_available(
 943        this: Entity<Self>,
 944        envelope: TypedEnvelope<proto::NewExternalAgentVersionAvailable>,
 945        mut cx: AsyncApp,
 946    ) -> Result<()> {
 947        this.update(&mut cx, |this, _| {
 948            if let Some(entry) = this.external_agents.get_mut(&*envelope.payload.name)
 949                && let Some(mut tx) = entry.server.take_new_version_available_tx()
 950            {
 951                tx.send(Some(envelope.payload.version)).ok();
 952                entry.server.set_new_version_available_tx(tx);
 953            }
 954        });
 955        Ok(())
 956    }
 957
 958    pub fn get_extension_id_for_agent(&self, name: &AgentId) -> Option<Arc<str>> {
 959        self.external_agents.get(name).and_then(|entry| {
 960            entry
 961                .server
 962                .as_any()
 963                .downcast_ref::<LocalExtensionArchiveAgent>()
 964                .map(|ext_agent| ext_agent.extension_id.clone())
 965        })
 966    }
 967}
 968
 969struct RemoteExternalAgentServer {
 970    project_id: u64,
 971    upstream_client: Entity<RemoteClient>,
 972    worktree_store: Entity<WorktreeStore>,
 973    name: AgentId,
 974    new_version_available_tx: Option<watch::Sender<Option<String>>>,
 975}
 976
 977impl ExternalAgentServer for RemoteExternalAgentServer {
 978    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
 979        self.new_version_available_tx.take()
 980    }
 981
 982    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
 983        self.new_version_available_tx = Some(tx);
 984    }
 985
 986    fn get_command(
 987        &self,
 988        extra_args: Vec<String>,
 989        extra_env: HashMap<String, String>,
 990        cx: &mut AsyncApp,
 991    ) -> Task<Result<AgentServerCommand>> {
 992        let project_id = self.project_id;
 993        let name = self.name.to_string();
 994        let upstream_client = self.upstream_client.downgrade();
 995        let worktree_store = self.worktree_store.clone();
 996        cx.spawn(async move |cx| {
 997            let root_dir = worktree_store.read_with(cx, |worktree_store, cx| {
 998                crate::Project::default_visible_worktree_paths(worktree_store, cx)
 999                    .into_iter()
1000                    .next()
1001                    .map(|path| path.display().to_string())
1002            });
1003
1004            let mut response = upstream_client
1005                .update(cx, |upstream_client, _| {
1006                    upstream_client
1007                        .proto_client()
1008                        .request(proto::GetAgentServerCommand {
1009                            project_id,
1010                            name,
1011                            root_dir,
1012                        })
1013                })?
1014                .await?;
1015            response.args.extend(extra_args);
1016            response.env.extend(extra_env);
1017
1018            Ok(AgentServerCommand {
1019                path: response.path.into(),
1020                args: response.args,
1021                env: Some(response.env.into_iter().collect()),
1022            })
1023        })
1024    }
1025
1026    fn as_any(&self) -> &dyn Any {
1027        self
1028    }
1029
1030    fn as_any_mut(&mut self) -> &mut dyn Any {
1031        self
1032    }
1033}
1034
1035fn asset_kind_for_archive_url(archive_url: &str) -> Result<AssetKind> {
1036    let archive_path = Url::parse(archive_url)
1037        .ok()
1038        .map(|url| url.path().to_string())
1039        .unwrap_or_else(|| archive_url.to_string());
1040
1041    if archive_path.ends_with(".zip") {
1042        Ok(AssetKind::Zip)
1043    } else if archive_path.ends_with(".tar.gz") || archive_path.ends_with(".tgz") {
1044        Ok(AssetKind::TarGz)
1045    } else if archive_path.ends_with(".tar.bz2") || archive_path.ends_with(".tbz2") {
1046        Ok(AssetKind::TarBz2)
1047    } else {
1048        bail!("unsupported archive type in URL: {archive_url}");
1049    }
1050}
1051
1052struct GithubReleaseArchive {
1053    repo_name_with_owner: String,
1054    tag: String,
1055    asset_name: String,
1056}
1057
1058fn github_release_archive_from_url(archive_url: &str) -> Option<GithubReleaseArchive> {
1059    fn decode_path_segment(segment: &str) -> Option<String> {
1060        percent_decode_str(segment)
1061            .decode_utf8()
1062            .ok()
1063            .map(|segment| segment.into_owned())
1064    }
1065
1066    let url = Url::parse(archive_url).ok()?;
1067    if url.scheme() != "https" || url.host_str()? != "github.com" {
1068        return None;
1069    }
1070
1071    let segments = url.path_segments()?.collect::<Vec<_>>();
1072    if segments.len() < 6 || segments[2] != "releases" || segments[3] != "download" {
1073        return None;
1074    }
1075
1076    Some(GithubReleaseArchive {
1077        repo_name_with_owner: format!("{}/{}", segments[0], segments[1]),
1078        tag: decode_path_segment(segments[4])?,
1079        asset_name: segments[5..]
1080            .iter()
1081            .map(|segment| decode_path_segment(segment))
1082            .collect::<Option<Vec<_>>>()?
1083            .join("/"),
1084    })
1085}
1086
1087fn sanitized_version_component(version: &str) -> String {
1088    let sanitized = version
1089        .chars()
1090        .map(|character| match character {
1091            'a'..='z' | 'A'..='Z' | '0'..='9' | '.' | '_' | '-' => character,
1092            _ => '-',
1093        })
1094        .collect::<String>();
1095
1096    if sanitized.is_empty() {
1097        "unknown".to_string()
1098    } else {
1099        sanitized
1100    }
1101}
1102
1103fn versioned_archive_cache_dir(
1104    base_dir: &Path,
1105    version: Option<&str>,
1106    archive_url: &str,
1107) -> PathBuf {
1108    let version = version.unwrap_or_default();
1109    let sanitized_version = sanitized_version_component(version);
1110
1111    let mut version_hasher = Sha256::new();
1112    version_hasher.update(version.as_bytes());
1113    let version_hash = format!("{:x}", version_hasher.finalize());
1114
1115    let mut url_hasher = Sha256::new();
1116    url_hasher.update(archive_url.as_bytes());
1117    let url_hash = format!("{:x}", url_hasher.finalize());
1118
1119    base_dir.join(format!(
1120        "v_{sanitized_version}_{}_{}",
1121        &version_hash[..16],
1122        &url_hash[..16],
1123    ))
1124}
1125
1126pub struct LocalExtensionArchiveAgent {
1127    pub fs: Arc<dyn Fs>,
1128    pub http_client: Arc<dyn HttpClient>,
1129    pub node_runtime: NodeRuntime,
1130    pub project_environment: Entity<ProjectEnvironment>,
1131    pub extension_id: Arc<str>,
1132    pub agent_id: Arc<str>,
1133    pub targets: HashMap<String, extension::TargetConfig>,
1134    pub env: HashMap<String, String>,
1135    pub version: Option<SharedString>,
1136    pub new_version_available_tx: Option<watch::Sender<Option<String>>>,
1137}
1138
1139impl ExternalAgentServer for LocalExtensionArchiveAgent {
1140    fn version(&self) -> Option<&SharedString> {
1141        self.version.as_ref()
1142    }
1143
1144    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1145        self.new_version_available_tx.take()
1146    }
1147
1148    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1149        self.new_version_available_tx = Some(tx);
1150    }
1151
1152    fn get_command(
1153        &self,
1154        extra_args: Vec<String>,
1155        extra_env: HashMap<String, String>,
1156        cx: &mut AsyncApp,
1157    ) -> Task<Result<AgentServerCommand>> {
1158        let fs = self.fs.clone();
1159        let http_client = self.http_client.clone();
1160        let node_runtime = self.node_runtime.clone();
1161        let project_environment = self.project_environment.downgrade();
1162        let extension_id = self.extension_id.clone();
1163        let agent_id = self.agent_id.clone();
1164        let targets = self.targets.clone();
1165        let base_env = self.env.clone();
1166        let version = self.version.clone();
1167
1168        cx.spawn(async move |cx| {
1169            // Get project environment
1170            let mut env = project_environment
1171                .update(cx, |project_environment, cx| {
1172                    project_environment.default_environment(cx)
1173                })?
1174                .await
1175                .unwrap_or_default();
1176
1177            // Merge manifest env and extra env
1178            env.extend(base_env);
1179            env.extend(extra_env);
1180
1181            let cache_key = format!("{}/{}", extension_id, agent_id);
1182            let dir = paths::external_agents_dir().join(&cache_key);
1183            fs.create_dir(&dir).await?;
1184
1185            // Determine platform key
1186            let os = if cfg!(target_os = "macos") {
1187                "darwin"
1188            } else if cfg!(target_os = "linux") {
1189                "linux"
1190            } else if cfg!(target_os = "windows") {
1191                "windows"
1192            } else {
1193                anyhow::bail!("unsupported OS");
1194            };
1195
1196            let arch = if cfg!(target_arch = "aarch64") {
1197                "aarch64"
1198            } else if cfg!(target_arch = "x86_64") {
1199                "x86_64"
1200            } else {
1201                anyhow::bail!("unsupported architecture");
1202            };
1203
1204            let platform_key = format!("{}-{}", os, arch);
1205            let target_config = targets.get(&platform_key).with_context(|| {
1206                format!(
1207                    "no target specified for platform '{}'. Available platforms: {}",
1208                    platform_key,
1209                    targets
1210                        .keys()
1211                        .map(|k| k.as_str())
1212                        .collect::<Vec<_>>()
1213                        .join(", ")
1214                )
1215            })?;
1216
1217            let archive_url = &target_config.archive;
1218            let version_dir = versioned_archive_cache_dir(
1219                &dir,
1220                version.as_ref().map(|version| version.as_ref()),
1221                archive_url,
1222            );
1223
1224            if !fs.is_dir(&version_dir).await {
1225                // Determine SHA256 for verification
1226                let sha256 = if let Some(provided_sha) = &target_config.sha256 {
1227                    // Use provided SHA256
1228                    Some(provided_sha.clone())
1229                } else if let Some(github_archive) = github_release_archive_from_url(archive_url) {
1230                    // Try to fetch SHA256 from GitHub API
1231                    if let Ok(release) = ::http_client::github::get_release_by_tag_name(
1232                        &github_archive.repo_name_with_owner,
1233                        &github_archive.tag,
1234                        http_client.clone(),
1235                    )
1236                    .await
1237                    {
1238                        // Find matching asset
1239                        if let Some(asset) = release
1240                            .assets
1241                            .iter()
1242                            .find(|a| a.name == github_archive.asset_name)
1243                        {
1244                            // Strip "sha256:" prefix if present
1245                            asset.digest.as_ref().map(|d| {
1246                                d.strip_prefix("sha256:")
1247                                    .map(|s| s.to_string())
1248                                    .unwrap_or_else(|| d.clone())
1249                            })
1250                        } else {
1251                            None
1252                        }
1253                    } else {
1254                        None
1255                    }
1256                } else {
1257                    None
1258                };
1259
1260                let asset_kind = asset_kind_for_archive_url(archive_url)?;
1261
1262                // Download and extract
1263                ::http_client::github_download::download_server_binary(
1264                    &*http_client,
1265                    archive_url,
1266                    sha256.as_deref(),
1267                    &version_dir,
1268                    asset_kind,
1269                )
1270                .await?;
1271            }
1272
1273            // Validate and resolve cmd path
1274            let cmd = &target_config.cmd;
1275
1276            let cmd_path = if cmd == "node" {
1277                // Use Zed's managed Node.js runtime
1278                node_runtime.binary_path().await?
1279            } else {
1280                if cmd.contains("..") {
1281                    anyhow::bail!("command path cannot contain '..': {}", cmd);
1282                }
1283
1284                if cmd.starts_with("./") || cmd.starts_with(".\\") {
1285                    // Relative to extraction directory
1286                    let cmd_path = version_dir.join(&cmd[2..]);
1287                    anyhow::ensure!(
1288                        fs.is_file(&cmd_path).await,
1289                        "Missing command {} after extraction",
1290                        cmd_path.to_string_lossy()
1291                    );
1292                    cmd_path
1293                } else {
1294                    // On PATH
1295                    anyhow::bail!("command must be relative (start with './'): {}", cmd);
1296                }
1297            };
1298
1299            let mut args = target_config.args.clone();
1300            args.extend(extra_args);
1301
1302            let command = AgentServerCommand {
1303                path: cmd_path,
1304                args,
1305                env: Some(env),
1306            };
1307
1308            Ok(command)
1309        })
1310    }
1311
1312    fn as_any(&self) -> &dyn Any {
1313        self
1314    }
1315
1316    fn as_any_mut(&mut self) -> &mut dyn Any {
1317        self
1318    }
1319}
1320
1321struct LocalRegistryArchiveAgent {
1322    fs: Arc<dyn Fs>,
1323    http_client: Arc<dyn HttpClient>,
1324    node_runtime: NodeRuntime,
1325    project_environment: Entity<ProjectEnvironment>,
1326    registry_id: Arc<str>,
1327    version: SharedString,
1328    targets: HashMap<String, RegistryTargetConfig>,
1329    env: HashMap<String, String>,
1330    new_version_available_tx: Option<watch::Sender<Option<String>>>,
1331}
1332
1333impl ExternalAgentServer for LocalRegistryArchiveAgent {
1334    fn version(&self) -> Option<&SharedString> {
1335        Some(&self.version)
1336    }
1337
1338    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1339        self.new_version_available_tx.take()
1340    }
1341
1342    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1343        self.new_version_available_tx = Some(tx);
1344    }
1345
1346    fn get_command(
1347        &self,
1348        extra_args: Vec<String>,
1349        extra_env: HashMap<String, String>,
1350        cx: &mut AsyncApp,
1351    ) -> Task<Result<AgentServerCommand>> {
1352        let fs = self.fs.clone();
1353        let http_client = self.http_client.clone();
1354        let node_runtime = self.node_runtime.clone();
1355        let project_environment = self.project_environment.downgrade();
1356        let registry_id = self.registry_id.clone();
1357        let targets = self.targets.clone();
1358        let settings_env = self.env.clone();
1359        let version = self.version.clone();
1360
1361        cx.spawn(async move |cx| {
1362            let mut env = project_environment
1363                .update(cx, |project_environment, cx| {
1364                    project_environment.default_environment(cx)
1365                })?
1366                .await
1367                .unwrap_or_default();
1368
1369            let dir = paths::external_agents_dir()
1370                .join("registry")
1371                .join(registry_id.as_ref());
1372            fs.create_dir(&dir).await?;
1373
1374            let os = if cfg!(target_os = "macos") {
1375                "darwin"
1376            } else if cfg!(target_os = "linux") {
1377                "linux"
1378            } else if cfg!(target_os = "windows") {
1379                "windows"
1380            } else {
1381                anyhow::bail!("unsupported OS");
1382            };
1383
1384            let arch = if cfg!(target_arch = "aarch64") {
1385                "aarch64"
1386            } else if cfg!(target_arch = "x86_64") {
1387                "x86_64"
1388            } else {
1389                anyhow::bail!("unsupported architecture");
1390            };
1391
1392            let platform_key = format!("{}-{}", os, arch);
1393            let target_config = targets.get(&platform_key).with_context(|| {
1394                format!(
1395                    "no target specified for platform '{}'. Available platforms: {}",
1396                    platform_key,
1397                    targets
1398                        .keys()
1399                        .map(|k| k.as_str())
1400                        .collect::<Vec<_>>()
1401                        .join(", ")
1402                )
1403            })?;
1404
1405            env.extend(target_config.env.clone());
1406            env.extend(extra_env);
1407            env.extend(settings_env);
1408
1409            let archive_url = &target_config.archive;
1410            let version_dir =
1411                versioned_archive_cache_dir(&dir, Some(version.as_ref()), archive_url);
1412
1413            if !fs.is_dir(&version_dir).await {
1414                let sha256 = if let Some(provided_sha) = &target_config.sha256 {
1415                    Some(provided_sha.clone())
1416                } else if let Some(github_archive) = github_release_archive_from_url(archive_url) {
1417                    if let Ok(release) = ::http_client::github::get_release_by_tag_name(
1418                        &github_archive.repo_name_with_owner,
1419                        &github_archive.tag,
1420                        http_client.clone(),
1421                    )
1422                    .await
1423                    {
1424                        if let Some(asset) = release
1425                            .assets
1426                            .iter()
1427                            .find(|a| a.name == github_archive.asset_name)
1428                        {
1429                            asset.digest.as_ref().and_then(|d| {
1430                                d.strip_prefix("sha256:")
1431                                    .map(|s| s.to_string())
1432                                    .or_else(|| Some(d.clone()))
1433                            })
1434                        } else {
1435                            None
1436                        }
1437                    } else {
1438                        None
1439                    }
1440                } else {
1441                    None
1442                };
1443
1444                let asset_kind = asset_kind_for_archive_url(archive_url)?;
1445
1446                ::http_client::github_download::download_server_binary(
1447                    &*http_client,
1448                    archive_url,
1449                    sha256.as_deref(),
1450                    &version_dir,
1451                    asset_kind,
1452                )
1453                .await?;
1454            }
1455
1456            let cmd = &target_config.cmd;
1457
1458            let cmd_path = if cmd == "node" {
1459                node_runtime.binary_path().await?
1460            } else {
1461                if cmd.contains("..") {
1462                    anyhow::bail!("command path cannot contain '..': {}", cmd);
1463                }
1464
1465                if cmd.starts_with("./") || cmd.starts_with(".\\") {
1466                    let cmd_path = version_dir.join(&cmd[2..]);
1467                    anyhow::ensure!(
1468                        fs.is_file(&cmd_path).await,
1469                        "Missing command {} after extraction",
1470                        cmd_path.to_string_lossy()
1471                    );
1472                    cmd_path
1473                } else {
1474                    anyhow::bail!("command must be relative (start with './'): {}", cmd);
1475                }
1476            };
1477
1478            let mut args = target_config.args.clone();
1479            args.extend(extra_args);
1480
1481            let command = AgentServerCommand {
1482                path: cmd_path,
1483                args,
1484                env: Some(env),
1485            };
1486
1487            Ok(command)
1488        })
1489    }
1490
1491    fn as_any(&self) -> &dyn Any {
1492        self
1493    }
1494
1495    fn as_any_mut(&mut self) -> &mut dyn Any {
1496        self
1497    }
1498}
1499
1500struct LocalRegistryNpxAgent {
1501    node_runtime: NodeRuntime,
1502    project_environment: Entity<ProjectEnvironment>,
1503    version: SharedString,
1504    package: SharedString,
1505    args: Vec<String>,
1506    distribution_env: HashMap<String, String>,
1507    settings_env: HashMap<String, String>,
1508    new_version_available_tx: Option<watch::Sender<Option<String>>>,
1509}
1510
1511impl ExternalAgentServer for LocalRegistryNpxAgent {
1512    fn version(&self) -> Option<&SharedString> {
1513        Some(&self.version)
1514    }
1515
1516    fn take_new_version_available_tx(&mut self) -> Option<watch::Sender<Option<String>>> {
1517        self.new_version_available_tx.take()
1518    }
1519
1520    fn set_new_version_available_tx(&mut self, tx: watch::Sender<Option<String>>) {
1521        self.new_version_available_tx = Some(tx);
1522    }
1523
1524    fn get_command(
1525        &self,
1526        extra_args: Vec<String>,
1527        extra_env: HashMap<String, String>,
1528        cx: &mut AsyncApp,
1529    ) -> Task<Result<AgentServerCommand>> {
1530        let node_runtime = self.node_runtime.clone();
1531        let project_environment = self.project_environment.downgrade();
1532        let package = self.package.clone();
1533        let args = self.args.clone();
1534        let distribution_env = self.distribution_env.clone();
1535        let settings_env = self.settings_env.clone();
1536
1537        cx.spawn(async move |cx| {
1538            let mut env = project_environment
1539                .update(cx, |project_environment, cx| {
1540                    project_environment.default_environment(cx)
1541                })?
1542                .await
1543                .unwrap_or_default();
1544
1545            let mut exec_args = vec!["--yes".to_string(), "--".to_string(), package.to_string()];
1546            exec_args.extend(args);
1547
1548            let npm_command = node_runtime
1549                .npm_command(
1550                    "exec",
1551                    &exec_args.iter().map(|a| a.as_str()).collect::<Vec<_>>(),
1552                )
1553                .await?;
1554
1555            env.extend(npm_command.env);
1556            env.extend(distribution_env);
1557            env.extend(extra_env);
1558            env.extend(settings_env);
1559
1560            let mut args = npm_command.args;
1561            args.extend(extra_args);
1562
1563            let command = AgentServerCommand {
1564                path: npm_command.path,
1565                args,
1566                env: Some(env),
1567            };
1568
1569            Ok(command)
1570        })
1571    }
1572
1573    fn as_any(&self) -> &dyn Any {
1574        self
1575    }
1576
1577    fn as_any_mut(&mut self) -> &mut dyn Any {
1578        self
1579    }
1580}
1581
1582struct LocalCustomAgent {
1583    project_environment: Entity<ProjectEnvironment>,
1584    command: AgentServerCommand,
1585}
1586
1587impl ExternalAgentServer for LocalCustomAgent {
1588    fn get_command(
1589        &self,
1590        extra_args: Vec<String>,
1591        extra_env: HashMap<String, String>,
1592        cx: &mut AsyncApp,
1593    ) -> Task<Result<AgentServerCommand>> {
1594        let mut command = self.command.clone();
1595        let project_environment = self.project_environment.downgrade();
1596        cx.spawn(async move |cx| {
1597            let mut env = project_environment
1598                .update(cx, |project_environment, cx| {
1599                    project_environment.default_environment(cx)
1600                })?
1601                .await
1602                .unwrap_or_default();
1603            env.extend(command.env.unwrap_or_default());
1604            env.extend(extra_env);
1605            command.env = Some(env);
1606            command.args.extend(extra_args);
1607            Ok(command)
1608        })
1609    }
1610
1611    fn as_any(&self) -> &dyn Any {
1612        self
1613    }
1614
1615    fn as_any_mut(&mut self) -> &mut dyn Any {
1616        self
1617    }
1618}
1619
1620#[derive(Default, Clone, JsonSchema, Debug, PartialEq, RegisterSetting)]
1621pub struct AllAgentServersSettings(pub HashMap<String, CustomAgentServerSettings>);
1622
1623impl std::ops::Deref for AllAgentServersSettings {
1624    type Target = HashMap<String, CustomAgentServerSettings>;
1625
1626    fn deref(&self) -> &Self::Target {
1627        &self.0
1628    }
1629}
1630
1631impl std::ops::DerefMut for AllAgentServersSettings {
1632    fn deref_mut(&mut self) -> &mut Self::Target {
1633        &mut self.0
1634    }
1635}
1636
1637impl AllAgentServersSettings {
1638    pub fn has_registry_agents(&self) -> bool {
1639        self.values()
1640            .any(|s| matches!(s, CustomAgentServerSettings::Registry { .. }))
1641    }
1642}
1643
1644#[derive(Clone, JsonSchema, Debug, PartialEq)]
1645pub enum CustomAgentServerSettings {
1646    Custom {
1647        command: AgentServerCommand,
1648        /// The default mode to use for this agent.
1649        ///
1650        /// Note: Not only all agents support modes.
1651        ///
1652        /// Default: None
1653        default_mode: Option<String>,
1654        /// The default model to use for this agent.
1655        ///
1656        /// This should be the model ID as reported by the agent.
1657        ///
1658        /// Default: None
1659        default_model: Option<String>,
1660        /// The favorite models for this agent.
1661        ///
1662        /// Default: []
1663        favorite_models: Vec<String>,
1664        /// Default values for session config options.
1665        ///
1666        /// This is a map from config option ID to value ID.
1667        ///
1668        /// Default: {}
1669        default_config_options: HashMap<String, String>,
1670        /// Favorited values for session config options.
1671        ///
1672        /// This is a map from config option ID to a list of favorited value IDs.
1673        ///
1674        /// Default: {}
1675        favorite_config_option_values: HashMap<String, Vec<String>>,
1676    },
1677    Extension {
1678        /// Additional environment variables to pass to the agent.
1679        ///
1680        /// Default: {}
1681        env: HashMap<String, String>,
1682        /// The default mode to use for this agent.
1683        ///
1684        /// Note: Not only all agents support modes.
1685        ///
1686        /// Default: None
1687        default_mode: Option<String>,
1688        /// The default model to use for this agent.
1689        ///
1690        /// This should be the model ID as reported by the agent.
1691        ///
1692        /// Default: None
1693        default_model: Option<String>,
1694        /// The favorite models for this agent.
1695        ///
1696        /// Default: []
1697        favorite_models: Vec<String>,
1698        /// Default values for session config options.
1699        ///
1700        /// This is a map from config option ID to value ID.
1701        ///
1702        /// Default: {}
1703        default_config_options: HashMap<String, String>,
1704        /// Favorited values for session config options.
1705        ///
1706        /// This is a map from config option ID to a list of favorited value IDs.
1707        ///
1708        /// Default: {}
1709        favorite_config_option_values: HashMap<String, Vec<String>>,
1710    },
1711    Registry {
1712        /// Additional environment variables to pass to the agent.
1713        ///
1714        /// Default: {}
1715        env: HashMap<String, String>,
1716        /// The default mode to use for this agent.
1717        ///
1718        /// Note: Not only all agents support modes.
1719        ///
1720        /// Default: None
1721        default_mode: Option<String>,
1722        /// The default model to use for this agent.
1723        ///
1724        /// This should be the model ID as reported by the agent.
1725        ///
1726        /// Default: None
1727        default_model: Option<String>,
1728        /// The favorite models for this agent.
1729        ///
1730        /// Default: []
1731        favorite_models: Vec<String>,
1732        /// Default values for session config options.
1733        ///
1734        /// This is a map from config option ID to value ID.
1735        ///
1736        /// Default: {}
1737        default_config_options: HashMap<String, String>,
1738        /// Favorited values for session config options.
1739        ///
1740        /// This is a map from config option ID to a list of favorited value IDs.
1741        ///
1742        /// Default: {}
1743        favorite_config_option_values: HashMap<String, Vec<String>>,
1744    },
1745}
1746
1747impl CustomAgentServerSettings {
1748    pub fn command(&self) -> Option<&AgentServerCommand> {
1749        match self {
1750            CustomAgentServerSettings::Custom { command, .. } => Some(command),
1751            CustomAgentServerSettings::Extension { .. }
1752            | CustomAgentServerSettings::Registry { .. } => None,
1753        }
1754    }
1755
1756    pub fn default_mode(&self) -> Option<&str> {
1757        match self {
1758            CustomAgentServerSettings::Custom { default_mode, .. }
1759            | CustomAgentServerSettings::Extension { default_mode, .. }
1760            | CustomAgentServerSettings::Registry { default_mode, .. } => default_mode.as_deref(),
1761        }
1762    }
1763
1764    pub fn default_model(&self) -> Option<&str> {
1765        match self {
1766            CustomAgentServerSettings::Custom { default_model, .. }
1767            | CustomAgentServerSettings::Extension { default_model, .. }
1768            | CustomAgentServerSettings::Registry { default_model, .. } => default_model.as_deref(),
1769        }
1770    }
1771
1772    pub fn favorite_models(&self) -> &[String] {
1773        match self {
1774            CustomAgentServerSettings::Custom {
1775                favorite_models, ..
1776            }
1777            | CustomAgentServerSettings::Extension {
1778                favorite_models, ..
1779            }
1780            | CustomAgentServerSettings::Registry {
1781                favorite_models, ..
1782            } => favorite_models,
1783        }
1784    }
1785
1786    pub fn default_config_option(&self, config_id: &str) -> Option<&str> {
1787        match self {
1788            CustomAgentServerSettings::Custom {
1789                default_config_options,
1790                ..
1791            }
1792            | CustomAgentServerSettings::Extension {
1793                default_config_options,
1794                ..
1795            }
1796            | CustomAgentServerSettings::Registry {
1797                default_config_options,
1798                ..
1799            } => default_config_options.get(config_id).map(|s| s.as_str()),
1800        }
1801    }
1802
1803    pub fn favorite_config_option_values(&self, config_id: &str) -> Option<&[String]> {
1804        match self {
1805            CustomAgentServerSettings::Custom {
1806                favorite_config_option_values,
1807                ..
1808            }
1809            | CustomAgentServerSettings::Extension {
1810                favorite_config_option_values,
1811                ..
1812            }
1813            | CustomAgentServerSettings::Registry {
1814                favorite_config_option_values,
1815                ..
1816            } => favorite_config_option_values
1817                .get(config_id)
1818                .map(|v| v.as_slice()),
1819        }
1820    }
1821}
1822
1823impl From<settings::CustomAgentServerSettings> for CustomAgentServerSettings {
1824    fn from(value: settings::CustomAgentServerSettings) -> Self {
1825        match value {
1826            settings::CustomAgentServerSettings::Custom {
1827                path,
1828                args,
1829                env,
1830                default_mode,
1831                default_model,
1832                favorite_models,
1833                default_config_options,
1834                favorite_config_option_values,
1835            } => CustomAgentServerSettings::Custom {
1836                command: AgentServerCommand {
1837                    path: PathBuf::from(shellexpand::tilde(&path.to_string_lossy()).as_ref()),
1838                    args,
1839                    env: Some(env),
1840                },
1841                default_mode,
1842                default_model,
1843                favorite_models,
1844                default_config_options,
1845                favorite_config_option_values,
1846            },
1847            settings::CustomAgentServerSettings::Extension {
1848                env,
1849                default_mode,
1850                default_model,
1851                default_config_options,
1852                favorite_models,
1853                favorite_config_option_values,
1854            } => CustomAgentServerSettings::Extension {
1855                env,
1856                default_mode,
1857                default_model,
1858                default_config_options,
1859                favorite_models,
1860                favorite_config_option_values,
1861            },
1862            settings::CustomAgentServerSettings::Registry {
1863                env,
1864                default_mode,
1865                default_model,
1866                default_config_options,
1867                favorite_models,
1868                favorite_config_option_values,
1869            } => CustomAgentServerSettings::Registry {
1870                env,
1871                default_mode,
1872                default_model,
1873                default_config_options,
1874                favorite_models,
1875                favorite_config_option_values,
1876            },
1877        }
1878    }
1879}
1880
1881impl settings::Settings for AllAgentServersSettings {
1882    fn from_settings(content: &settings::SettingsContent) -> Self {
1883        let agent_settings = content.agent_servers.clone().unwrap();
1884        Self(
1885            agent_settings
1886                .0
1887                .into_iter()
1888                .map(|(k, v)| (k, v.into()))
1889                .collect(),
1890        )
1891    }
1892}
1893
1894#[cfg(test)]
1895mod tests {
1896    use super::*;
1897    use crate::agent_registry_store::{
1898        AgentRegistryStore, RegistryAgent, RegistryAgentMetadata, RegistryNpxAgent,
1899    };
1900    use crate::worktree_store::{WorktreeIdCounter, WorktreeStore};
1901    use gpui::{AppContext as _, TestAppContext};
1902    use node_runtime::NodeRuntime;
1903    use settings::Settings as _;
1904
1905    fn make_npx_agent(id: &str, version: &str) -> RegistryAgent {
1906        let id = SharedString::from(id.to_string());
1907        RegistryAgent::Npx(RegistryNpxAgent {
1908            metadata: RegistryAgentMetadata {
1909                id: AgentId::new(id.clone()),
1910                name: id.clone(),
1911                description: SharedString::from(""),
1912                version: SharedString::from(version.to_string()),
1913                repository: None,
1914                website: None,
1915                icon_path: None,
1916            },
1917            package: id,
1918            args: Vec::new(),
1919            env: HashMap::default(),
1920        })
1921    }
1922
1923    fn init_test_settings(cx: &mut TestAppContext) {
1924        cx.update(|cx| {
1925            let settings_store = SettingsStore::test(cx);
1926            cx.set_global(settings_store);
1927        });
1928    }
1929
1930    fn init_registry(
1931        cx: &mut TestAppContext,
1932        agents: Vec<RegistryAgent>,
1933    ) -> gpui::Entity<AgentRegistryStore> {
1934        cx.update(|cx| AgentRegistryStore::init_test_global(cx, agents))
1935    }
1936
1937    fn set_registry_settings(cx: &mut TestAppContext, agent_names: &[&str]) {
1938        cx.update(|cx| {
1939            AllAgentServersSettings::override_global(
1940                AllAgentServersSettings(
1941                    agent_names
1942                        .iter()
1943                        .map(|name| {
1944                            (
1945                                name.to_string(),
1946                                settings::CustomAgentServerSettings::Registry {
1947                                    env: HashMap::default(),
1948                                    default_mode: None,
1949                                    default_model: None,
1950                                    favorite_models: Vec::new(),
1951                                    default_config_options: HashMap::default(),
1952                                    favorite_config_option_values: HashMap::default(),
1953                                }
1954                                .into(),
1955                            )
1956                        })
1957                        .collect(),
1958                ),
1959                cx,
1960            );
1961        });
1962    }
1963
1964    fn create_agent_server_store(cx: &mut TestAppContext) -> gpui::Entity<AgentServerStore> {
1965        cx.update(|cx| {
1966            let fs: Arc<dyn Fs> = fs::FakeFs::new(cx.background_executor().clone());
1967            let worktree_store =
1968                cx.new(|cx| WorktreeStore::local(false, fs.clone(), WorktreeIdCounter::get(cx)));
1969            let project_environment = cx.new(|cx| {
1970                crate::ProjectEnvironment::new(None, worktree_store.downgrade(), None, false, cx)
1971            });
1972            let http_client = http_client::FakeHttpClient::with_404_response();
1973
1974            cx.new(|cx| {
1975                AgentServerStore::local(
1976                    NodeRuntime::unavailable(),
1977                    fs,
1978                    project_environment,
1979                    http_client,
1980                    cx,
1981                )
1982            })
1983        })
1984    }
1985
1986    #[test]
1987    fn detects_supported_archive_suffixes() {
1988        assert!(matches!(
1989            asset_kind_for_archive_url("https://example.com/agent.zip"),
1990            Ok(AssetKind::Zip)
1991        ));
1992        assert!(matches!(
1993            asset_kind_for_archive_url("https://example.com/agent.zip?download=1"),
1994            Ok(AssetKind::Zip)
1995        ));
1996        assert!(matches!(
1997            asset_kind_for_archive_url("https://example.com/agent.tar.gz"),
1998            Ok(AssetKind::TarGz)
1999        ));
2000        assert!(matches!(
2001            asset_kind_for_archive_url("https://example.com/agent.tar.gz?download=1#latest"),
2002            Ok(AssetKind::TarGz)
2003        ));
2004        assert!(matches!(
2005            asset_kind_for_archive_url("https://example.com/agent.tgz"),
2006            Ok(AssetKind::TarGz)
2007        ));
2008        assert!(matches!(
2009            asset_kind_for_archive_url("https://example.com/agent.tgz#download"),
2010            Ok(AssetKind::TarGz)
2011        ));
2012        assert!(matches!(
2013            asset_kind_for_archive_url("https://example.com/agent.tar.bz2"),
2014            Ok(AssetKind::TarBz2)
2015        ));
2016        assert!(matches!(
2017            asset_kind_for_archive_url("https://example.com/agent.tar.bz2?download=1"),
2018            Ok(AssetKind::TarBz2)
2019        ));
2020        assert!(matches!(
2021            asset_kind_for_archive_url("https://example.com/agent.tbz2"),
2022            Ok(AssetKind::TarBz2)
2023        ));
2024        assert!(matches!(
2025            asset_kind_for_archive_url("https://example.com/agent.tbz2#download"),
2026            Ok(AssetKind::TarBz2)
2027        ));
2028    }
2029
2030    #[test]
2031    fn parses_github_release_archive_urls() {
2032        let github_archive = github_release_archive_from_url(
2033            "https://github.com/owner/repo/releases/download/release%2F2.3.5/agent.tar.bz2?download=1",
2034        )
2035        .unwrap();
2036
2037        assert_eq!(github_archive.repo_name_with_owner, "owner/repo");
2038        assert_eq!(github_archive.tag, "release/2.3.5");
2039        assert_eq!(github_archive.asset_name, "agent.tar.bz2");
2040    }
2041
2042    #[test]
2043    fn rejects_unsupported_archive_suffixes() {
2044        let error = asset_kind_for_archive_url("https://example.com/agent.tar.xz")
2045            .err()
2046            .map(|error| error.to_string());
2047
2048        assert_eq!(
2049            error,
2050            Some("unsupported archive type in URL: https://example.com/agent.tar.xz".to_string()),
2051        );
2052    }
2053
2054    #[test]
2055    fn versioned_archive_cache_dir_includes_version_before_url_hash() {
2056        let slash_version_dir = versioned_archive_cache_dir(
2057            Path::new("/tmp/agents"),
2058            Some("release/2.3.5"),
2059            "https://example.com/agent.zip",
2060        );
2061        let colon_version_dir = versioned_archive_cache_dir(
2062            Path::new("/tmp/agents"),
2063            Some("release:2.3.5"),
2064            "https://example.com/agent.zip",
2065        );
2066        let file_name = slash_version_dir
2067            .file_name()
2068            .and_then(|name| name.to_str())
2069            .expect("cache directory should have a file name");
2070
2071        assert!(file_name.starts_with("v_release-2.3.5_"));
2072        assert_ne!(slash_version_dir, colon_version_dir);
2073    }
2074
2075    #[gpui::test]
2076    fn test_version_change_sends_notification(cx: &mut TestAppContext) {
2077        init_test_settings(cx);
2078        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2079        set_registry_settings(cx, &["test-agent"]);
2080        let store = create_agent_server_store(cx);
2081
2082        // Verify the agent was registered with version 1.0.0.
2083        store.read_with(cx, |store, _| {
2084            let entry = store
2085                .external_agents
2086                .get(&AgentId::new("test-agent"))
2087                .expect("agent should be registered");
2088            assert_eq!(
2089                entry.server.version().map(|v| v.to_string()),
2090                Some("1.0.0".to_string())
2091            );
2092        });
2093
2094        // Set up a watch channel and store the tx on the agent.
2095        let (tx, mut rx) = watch::channel::<Option<String>>(None);
2096        store.update(cx, |store, _| {
2097            let entry = store
2098                .external_agents
2099                .get_mut(&AgentId::new("test-agent"))
2100                .expect("agent should be registered");
2101            entry.server.set_new_version_available_tx(tx);
2102        });
2103
2104        // Update the registry to version 2.0.0.
2105        registry.update(cx, |store, cx| {
2106            store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx);
2107        });
2108        cx.run_until_parked();
2109
2110        // The watch channel should have received the new version.
2111        assert_eq!(rx.borrow().as_deref(), Some("2.0.0"));
2112    }
2113
2114    #[gpui::test]
2115    fn test_same_version_preserves_tx(cx: &mut TestAppContext) {
2116        init_test_settings(cx);
2117        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2118        set_registry_settings(cx, &["test-agent"]);
2119        let store = create_agent_server_store(cx);
2120
2121        let (tx, mut rx) = watch::channel::<Option<String>>(None);
2122        store.update(cx, |store, _| {
2123            let entry = store
2124                .external_agents
2125                .get_mut(&AgentId::new("test-agent"))
2126                .expect("agent should be registered");
2127            entry.server.set_new_version_available_tx(tx);
2128        });
2129
2130        // "Refresh" the registry with the same version.
2131        registry.update(cx, |store, cx| {
2132            store.set_agents(vec![make_npx_agent("test-agent", "1.0.0")], cx);
2133        });
2134        cx.run_until_parked();
2135
2136        // No notification should have been sent.
2137        assert_eq!(rx.borrow().as_deref(), None);
2138
2139        // The tx should have been transferred to the rebuilt agent entry.
2140        store.update(cx, |store, _| {
2141            let entry = store
2142                .external_agents
2143                .get_mut(&AgentId::new("test-agent"))
2144                .expect("agent should be registered");
2145            assert!(
2146                entry.server.take_new_version_available_tx().is_some(),
2147                "tx should have been transferred to the rebuilt agent"
2148            );
2149        });
2150    }
2151
2152    #[gpui::test]
2153    fn test_no_tx_stored_does_not_panic_on_version_change(cx: &mut TestAppContext) {
2154        init_test_settings(cx);
2155        let registry = init_registry(cx, vec![make_npx_agent("test-agent", "1.0.0")]);
2156        set_registry_settings(cx, &["test-agent"]);
2157        let _store = create_agent_server_store(cx);
2158
2159        // Update the registry without having stored any tx — should not panic.
2160        registry.update(cx, |store, cx| {
2161            store.set_agents(vec![make_npx_agent("test-agent", "2.0.0")], cx);
2162        });
2163        cx.run_until_parked();
2164    }
2165
2166    #[gpui::test]
2167    fn test_multiple_agents_independent_notifications(cx: &mut TestAppContext) {
2168        init_test_settings(cx);
2169        let registry = init_registry(
2170            cx,
2171            vec![
2172                make_npx_agent("agent-a", "1.0.0"),
2173                make_npx_agent("agent-b", "3.0.0"),
2174            ],
2175        );
2176        set_registry_settings(cx, &["agent-a", "agent-b"]);
2177        let store = create_agent_server_store(cx);
2178
2179        let (tx_a, mut rx_a) = watch::channel::<Option<String>>(None);
2180        let (tx_b, mut rx_b) = watch::channel::<Option<String>>(None);
2181        store.update(cx, |store, _| {
2182            store
2183                .external_agents
2184                .get_mut(&AgentId::new("agent-a"))
2185                .expect("agent-a should be registered")
2186                .server
2187                .set_new_version_available_tx(tx_a);
2188            store
2189                .external_agents
2190                .get_mut(&AgentId::new("agent-b"))
2191                .expect("agent-b should be registered")
2192                .server
2193                .set_new_version_available_tx(tx_b);
2194        });
2195
2196        // Update only agent-a to a new version; agent-b stays the same.
2197        registry.update(cx, |store, cx| {
2198            store.set_agents(
2199                vec![
2200                    make_npx_agent("agent-a", "2.0.0"),
2201                    make_npx_agent("agent-b", "3.0.0"),
2202                ],
2203                cx,
2204            );
2205        });
2206        cx.run_until_parked();
2207
2208        // agent-a should have received a notification.
2209        assert_eq!(rx_a.borrow().as_deref(), Some("2.0.0"));
2210
2211        // agent-b should NOT have received a notification.
2212        assert_eq!(rx_b.borrow().as_deref(), None);
2213
2214        // agent-b's tx should have been transferred.
2215        store.update(cx, |store, _| {
2216            assert!(
2217                store
2218                    .external_agents
2219                    .get_mut(&AgentId::new("agent-b"))
2220                    .expect("agent-b should be registered")
2221                    .server
2222                    .take_new_version_available_tx()
2223                    .is_some(),
2224                "agent-b tx should have been transferred"
2225            );
2226        });
2227    }
2228}