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