agent: Allow env overrides for extension and registry agents (#47275)

Ben Brandt created

Align the behavior with built-in and custom agents

Release Notes:

- N/A

Change summary

crates/agent_servers/src/custom.rs         |  5 ++
crates/agent_ui/src/agent_configuration.rs |  2 
crates/agent_ui/src/agent_registry_ui.rs   |  1 
crates/project/src/agent_server_store.rs   | 53 +++++++++++++++++++++--
crates/settings_content/src/agent.rs       | 14 +++++
5 files changed, 68 insertions(+), 7 deletions(-)

Detailed changes

crates/agent_servers/src/custom.rs 🔗

@@ -91,6 +91,7 @@ impl AgentServer for CustomAgentServer {
                 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                     default_model: None,
                     default_mode: None,
+                    env: Default::default(),
                     favorite_models: Vec::new(),
                     default_config_options: Default::default(),
                     favorite_config_option_values: Default::default(),
@@ -139,6 +140,7 @@ impl AgentServer for CustomAgentServer {
                 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                     default_model: None,
                     default_mode: None,
+                    env: Default::default(),
                     favorite_models: Vec::new(),
                     default_config_options: Default::default(),
                     favorite_config_option_values: Default::default(),
@@ -179,6 +181,7 @@ impl AgentServer for CustomAgentServer {
                 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                     default_model: None,
                     default_mode: None,
+                    env: Default::default(),
                     favorite_models: Vec::new(),
                     default_config_options: Default::default(),
                     favorite_config_option_values: Default::default(),
@@ -231,6 +234,7 @@ impl AgentServer for CustomAgentServer {
                 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                     default_model: None,
                     default_mode: None,
+                    env: Default::default(),
                     favorite_models: Vec::new(),
                     default_config_options: Default::default(),
                     favorite_config_option_values: Default::default(),
@@ -292,6 +296,7 @@ impl AgentServer for CustomAgentServer {
                 .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                     default_model: None,
                     default_mode: None,
+                    env: Default::default(),
                     favorite_models: Vec::new(),
                     default_config_options: Default::default(),
                     favorite_config_option_values: Default::default(),

crates/agent_ui/src/agent_configuration.rs 🔗

@@ -1422,7 +1422,7 @@ async fn open_new_agent_servers_entry_in_settings_editor(
                             settings::CustomAgentServerSettings::Custom {
                                 path: "path_to_executable".into(),
                                 args: vec![],
-                                env: Some(HashMap::default()),
+                                env: HashMap::default(),
                                 default_mode: None,
                                 default_model: None,
                                 favorite_models: vec![],

crates/agent_ui/src/agent_registry_ui.rs 🔗

@@ -475,6 +475,7 @@ impl AgentRegistryPage {
                                 settings::CustomAgentServerSettings::Registry {
                                     default_mode: None,
                                     default_model: None,
+                                    env: Default::default(),
                                     favorite_models: Vec::new(),
                                     default_config_options: HashMap::default(),
                                     favorite_config_option_values: HashMap::default(),

crates/project/src/agent_server_store.rs 🔗

@@ -466,11 +466,26 @@ impl AgentServerStore {
                         let icon = icon_path;
                         let agent_server_name = ExternalAgentServerName(agent_name.clone().into());
                         self.external_agents
-                            .entry(agent_server_name)
+                            .entry(agent_server_name.clone())
                             .and_modify(|entry| {
                                 entry.icon = icon_shared.clone();
                                 entry.display_name = Some(display_name.clone());
                                 entry.source = ExternalAgentSource::Extension;
+                            })
+                            .or_insert_with(|| {
+                                ExternalAgentEntry::new(
+                                    Box::new(RemoteExternalAgentServer {
+                                        project_id: *project_id,
+                                        upstream_client: upstream_client.clone(),
+                                        name: agent_server_name.clone(),
+                                        status_tx: None,
+                                        new_version_available_tx: None,
+                                    })
+                                        as Box<dyn ExternalAgentServer>,
+                                    ExternalAgentSource::Extension,
+                                    icon_shared.clone(),
+                                    Some(display_name.clone()),
+                                )
                             });
 
                         agents.push(ExternalExtensionAgent {
@@ -725,7 +740,7 @@ impl AgentServerStore {
                         ),
                     );
                 }
-                CustomAgentServerSettings::Registry { .. } => {
+                CustomAgentServerSettings::Registry { env, .. } => {
                     let Some(agent) = registry_agents_by_id.get(name) else {
                         if registry_store.is_some() {
                             log::warn!("Registry agent '{}' not found in ACP registry", name);
@@ -751,6 +766,7 @@ impl AgentServerStore {
                                 project_environment: project_environment.clone(),
                                 registry_id: Arc::from(name.as_str()),
                                 targets: agent.targets.clone(),
+                                env: env.clone(),
                             }) as Box<dyn ExternalAgentServer>,
                             ExternalAgentSource::Registry,
                             agent.icon_path.clone(),
@@ -764,6 +780,18 @@ impl AgentServerStore {
 
         for (agent_name, ext_id, targets, env, icon_path, display_name) in extension_agents.iter() {
             let name = ExternalAgentServerName(agent_name.clone().into());
+            let mut env = env.clone();
+            if let Some(settings_env) =
+                new_settings
+                    .custom
+                    .get(agent_name.as_ref())
+                    .and_then(|settings| match settings {
+                        CustomAgentServerSettings::Extension { env, .. } => Some(env.clone()),
+                        _ => None,
+                    })
+            {
+                env.extend(settings_env);
+            }
             let icon = icon_path
                 .as_ref()
                 .map(|path| SharedString::from(path.clone()));
@@ -778,7 +806,7 @@ impl AgentServerStore {
                         project_environment: project_environment.clone(),
                         extension_id: Arc::from(&**ext_id),
                         targets: targets.clone(),
-                        env: env.clone(),
+                        env,
                         agent_id: agent_name.clone(),
                     }) as Box<dyn ExternalAgentServer>,
                     ExternalAgentSource::Extension,
@@ -2099,6 +2127,7 @@ struct LocalRegistryArchiveAgent {
     project_environment: Entity<ProjectEnvironment>,
     registry_id: Arc<str>,
     targets: HashMap<String, RegistryTargetConfig>,
+    env: HashMap<String, String>,
 }
 
 impl ExternalAgentServer for LocalRegistryArchiveAgent {
@@ -2116,6 +2145,7 @@ impl ExternalAgentServer for LocalRegistryArchiveAgent {
         let project_environment = self.project_environment.downgrade();
         let registry_id = self.registry_id.clone();
         let targets = self.targets.clone();
+        let settings_env = self.env.clone();
 
         let root_dir: Arc<Path> = root_dir
             .map(|root_dir| Path::new(root_dir))
@@ -2172,6 +2202,7 @@ impl ExternalAgentServer for LocalRegistryArchiveAgent {
 
             env.extend(target_config.env.clone());
             env.extend(extra_env);
+            env.extend(settings_env);
 
             let archive_url = &target_config.archive;
 
@@ -2420,6 +2451,10 @@ pub enum CustomAgentServerSettings {
         favorite_config_option_values: HashMap<String, Vec<String>>,
     },
     Extension {
+        /// Additional environment variables to pass to the agent.
+        ///
+        /// Default: {}
+        env: HashMap<String, String>,
         /// The default mode to use for this agent.
         ///
         /// Note: Not only all agents support modes.
@@ -2450,6 +2485,10 @@ pub enum CustomAgentServerSettings {
         favorite_config_option_values: HashMap<String, Vec<String>>,
     },
     Registry {
+        /// Additional environment variables to pass to the agent.
+        ///
+        /// Default: {}
+        env: HashMap<String, String>,
         /// The default mode to use for this agent.
         ///
         /// Note: Not only all agents support modes.
@@ -2573,7 +2612,7 @@ impl From<settings::CustomAgentServerSettings> for CustomAgentServerSettings {
                 command: AgentServerCommand {
                     path: PathBuf::from(shellexpand::tilde(&path.to_string_lossy()).as_ref()),
                     args,
-                    env,
+                    env: Some(env),
                 },
                 default_mode,
                 default_model,
@@ -2582,12 +2621,14 @@ impl From<settings::CustomAgentServerSettings> for CustomAgentServerSettings {
                 favorite_config_option_values,
             },
             settings::CustomAgentServerSettings::Extension {
+                env,
                 default_mode,
                 default_model,
                 default_config_options,
                 favorite_models,
                 favorite_config_option_values,
             } => CustomAgentServerSettings::Extension {
+                env,
                 default_mode,
                 default_model,
                 default_config_options,
@@ -2595,12 +2636,14 @@ impl From<settings::CustomAgentServerSettings> for CustomAgentServerSettings {
                 favorite_config_option_values,
             },
             settings::CustomAgentServerSettings::Registry {
+                env,
                 default_mode,
                 default_model,
                 default_config_options,
                 favorite_models,
                 favorite_config_option_values,
             } => CustomAgentServerSettings::Registry {
+                env,
                 default_mode,
                 default_model,
                 default_config_options,
@@ -2953,7 +2996,7 @@ mod extension_agent_tests {
         let settings = settings::CustomAgentServerSettings::Custom {
             path: PathBuf::from("~/custom/agent"),
             args: vec!["serve".into()],
-            env: None,
+            env: Default::default(),
             default_mode: None,
             default_model: None,
             favorite_models: vec![],

crates/settings_content/src/agent.rs 🔗

@@ -432,7 +432,9 @@ pub enum CustomAgentServerSettings {
         path: PathBuf,
         #[serde(default)]
         args: Vec<String>,
-        env: Option<HashMap<String, String>>,
+        /// Default: {}
+        #[serde(default)]
+        env: HashMap<String, String>,
         /// The default mode to use for this agent.
         ///
         /// Note: Not only all agents support modes.
@@ -468,6 +470,11 @@ pub enum CustomAgentServerSettings {
         favorite_config_option_values: HashMap<String, Vec<String>>,
     },
     Extension {
+        /// Additional environment variables to pass to the agent.
+        ///
+        /// Default: {}
+        #[serde(default)]
+        env: HashMap<String, String>,
         /// The default mode to use for this agent.
         ///
         /// Note: Not only all agents support modes.
@@ -503,6 +510,11 @@ pub enum CustomAgentServerSettings {
         favorite_config_option_values: HashMap<String, Vec<String>>,
     },
     Registry {
+        /// Additional environment variables to pass to the agent.
+        ///
+        /// Default: {}
+        #[serde(default)]
+        env: HashMap<String, String>,
         /// The default mode to use for this agent.
         ///
         /// Note: Not only all agents support modes.