2443 lines · 88.5 KB · Rust
RawDownload
use std::{    any::Any,    borrow::Borrow,    path::{Path, PathBuf},    str::FromStr as _,    sync::Arc,    time::Duration,}; use anyhow::{Context as _, Result, bail};use collections::HashMap;use fs::{Fs, RemoveOptions, RenameOptions};use futures::StreamExt as _;use gpui::{    AppContext as _, AsyncApp, Context, Entity, EventEmitter, SharedString, Subscription, Task,};use http_client::{HttpClient, github::AssetKind};use node_runtime::NodeRuntime;use remote::RemoteClient;use rpc::{    AnyProtoClient, TypedEnvelope,    proto::{self, ExternalExtensionAgent},};use schemars::JsonSchema;use semver::Version;use serde::{Deserialize, Serialize};use settings::{RegisterSetting, SettingsStore};use task::{Shell, SpawnInTerminal};use util::{ResultExt as _, debug_panic}; use crate::ProjectEnvironment; #[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema)]pub struct AgentServerCommand {    #[serde(rename = "command")]    pub path: PathBuf,    #[serde(default)]    pub args: Vec<String>,    pub env: Option<HashMap<String, String>>,} impl std::fmt::Debug for AgentServerCommand {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        let filtered_env = self.env.as_ref().map(|env| {            env.iter()                .map(|(k, v)| {                    (                        k,                        if util::redact::should_redact(k) {                            "[REDACTED]"                        } else {                            v                        },                    )                })                .collect::<Vec<_>>()        });         f.debug_struct("AgentServerCommand")            .field("path", &self.path)            .field("args", &self.args)            .field("env", &filtered_env)            .finish()    }} #[derive(Clone, Debug, PartialEq, Eq, Hash)]pub struct ExternalAgentServerName(pub SharedString); impl std::fmt::Display for ExternalAgentServerName {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}", self.0)    }} impl From<&'static str> for ExternalAgentServerName {    fn from(value: &'static str) -> Self {        ExternalAgentServerName(value.into())    }} impl From<ExternalAgentServerName> for SharedString {    fn from(value: ExternalAgentServerName) -> Self {        value.0    }} impl Borrow<str> for ExternalAgentServerName {    fn borrow(&self) -> &str {        &self.0    }} pub trait ExternalAgentServer {    fn get_command(        &mut self,        root_dir: Option<&str>,        extra_env: HashMap<String, String>,        status_tx: Option<watch::Sender<SharedString>>,        new_version_available_tx: Option<watch::Sender<Option<String>>>,        cx: &mut AsyncApp,    ) -> Task<Result<(AgentServerCommand, String, Option<task::SpawnInTerminal>)>>;     fn as_any_mut(&mut self) -> &mut dyn Any;} impl dyn ExternalAgentServer {    fn downcast_mut<T: ExternalAgentServer + 'static>(&mut self) -> Option<&mut T> {        self.as_any_mut().downcast_mut()    }} enum AgentServerStoreState {    Local {        node_runtime: NodeRuntime,        fs: Arc<dyn Fs>,        project_environment: Entity<ProjectEnvironment>,        downstream_client: Option<(u64, AnyProtoClient)>,        settings: Option<AllAgentServersSettings>,        http_client: Arc<dyn HttpClient>,        extension_agents: Vec<(            Arc<str>,            String,            HashMap<String, extension::TargetConfig>,            HashMap<String, String>,            Option<String>,        )>,        _subscriptions: [Subscription; 1],    },    Remote {        project_id: u64,        upstream_client: Entity<RemoteClient>,    },    Collab,} pub struct AgentServerStore {    state: AgentServerStoreState,    external_agents: HashMap<ExternalAgentServerName, Box<dyn ExternalAgentServer>>,    agent_icons: HashMap<ExternalAgentServerName, SharedString>,    agent_display_names: HashMap<ExternalAgentServerName, SharedString>,} pub struct AgentServersUpdated; impl EventEmitter<AgentServersUpdated> for AgentServerStore {} #[cfg(test)]mod ext_agent_tests {    use super::*;    use std::{collections::HashSet, fmt::Write as _};     // Helper to build a store in Collab mode so we can mutate internal maps without    // needing to spin up a full project environment.    fn collab_store() -> AgentServerStore {        AgentServerStore {            state: AgentServerStoreState::Collab,            external_agents: HashMap::default(),            agent_icons: HashMap::default(),            agent_display_names: HashMap::default(),        }    }     // A simple fake that implements ExternalAgentServer without needing async plumbing.    struct NoopExternalAgent;     impl ExternalAgentServer for NoopExternalAgent {        fn get_command(            &mut self,            _root_dir: Option<&str>,            _extra_env: HashMap<String, String>,            _status_tx: Option<watch::Sender<SharedString>>,            _new_version_available_tx: Option<watch::Sender<Option<String>>>,            _cx: &mut AsyncApp,        ) -> Task<Result<(AgentServerCommand, String, Option<task::SpawnInTerminal>)>> {            Task::ready(Ok((                AgentServerCommand {                    path: PathBuf::from("noop"),                    args: Vec::new(),                    env: None,                },                "".to_string(),                None,            )))        }         fn as_any_mut(&mut self) -> &mut dyn Any {            self        }    }     #[test]    fn external_agent_server_name_display() {        let name = ExternalAgentServerName(SharedString::from("Ext: Tool"));        let mut s = String::new();        write!(&mut s, "{name}").unwrap();        assert_eq!(s, "Ext: Tool");    }     #[test]    fn sync_extension_agents_removes_previous_extension_entries() {        let mut store = collab_store();         // Seed with a couple of agents that will be replaced by extensions        store.external_agents.insert(            ExternalAgentServerName(SharedString::from("foo-agent")),            Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,        );        store.external_agents.insert(            ExternalAgentServerName(SharedString::from("bar-agent")),            Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,        );        store.external_agents.insert(            ExternalAgentServerName(SharedString::from("custom")),            Box::new(NoopExternalAgent) as Box<dyn ExternalAgentServer>,        );         // Simulate the removal phase: if we're syncing extensions that provide        // "foo-agent" and "bar-agent", those should be removed first        let extension_agent_names: HashSet<String> =            ["foo-agent".to_string(), "bar-agent".to_string()]                .into_iter()                .collect();         let keys_to_remove: Vec<_> = store            .external_agents            .keys()            .filter(|name| extension_agent_names.contains(name.0.as_ref()))            .cloned()            .collect();         for key in keys_to_remove {            store.external_agents.remove(&key);        }         // Only the custom entry should remain.        let remaining: Vec<_> = store            .external_agents            .keys()            .map(|k| k.0.to_string())            .collect();        assert_eq!(remaining, vec!["custom".to_string()]);    }} impl AgentServerStore {    /// Synchronizes extension-provided agent servers with the store.    pub fn sync_extension_agents<'a, I>(        &mut self,        manifests: I,        extensions_dir: PathBuf,        cx: &mut Context<Self>,    ) where        I: IntoIterator<Item = (&'a str, &'a extension::ExtensionManifest)>,    {        // Collect manifests first so we can iterate twice        let manifests: Vec<_> = manifests.into_iter().collect();         // Remove all extension-provided agents        // (They will be re-added below if they're in the currently installed extensions)        self.external_agents.retain(|name, agent| {            if agent.downcast_mut::<LocalExtensionArchiveAgent>().is_some() {                self.agent_icons.remove(name);                self.agent_display_names.remove(name);                false            } else {                // Keep the hardcoded external agents that don't come from extensions                // (In the future we may move these over to being extensions too.)                true            }        });         // Insert agent servers from extension manifests        match &mut self.state {            AgentServerStoreState::Local {                extension_agents, ..            } => {                extension_agents.clear();                for (ext_id, manifest) in manifests {                    for (agent_name, agent_entry) in &manifest.agent_servers {                        // Store absolute icon path if provided, resolving symlinks for dev extensions                        // Store display name from manifest                        self.agent_display_names.insert(                            ExternalAgentServerName(agent_name.clone().into()),                            SharedString::from(agent_entry.name.clone()),                        );                         let icon_path = if let Some(icon) = &agent_entry.icon {                            let icon_path = extensions_dir.join(ext_id).join(icon);                            // Canonicalize to resolve symlinks (dev extensions are symlinked)                            let absolute_icon_path = icon_path                                .canonicalize()                                .unwrap_or(icon_path)                                .to_string_lossy()                                .to_string();                            self.agent_icons.insert(                                ExternalAgentServerName(agent_name.clone().into()),                                SharedString::from(absolute_icon_path.clone()),                            );                            Some(absolute_icon_path)                        } else {                            None                        };                         extension_agents.push((                            agent_name.clone(),                            ext_id.to_owned(),                            agent_entry.targets.clone(),                            agent_entry.env.clone(),                            icon_path,                        ));                    }                }                self.reregister_agents(cx);            }            AgentServerStoreState::Remote {                project_id,                upstream_client,            } => {                let mut agents = vec![];                for (ext_id, manifest) in manifests {                    for (agent_name, agent_entry) in &manifest.agent_servers {                        // Store display name from manifest                        self.agent_display_names.insert(                            ExternalAgentServerName(agent_name.clone().into()),                            SharedString::from(agent_entry.name.clone()),                        );                         // Store absolute icon path if provided, resolving symlinks for dev extensions                        let icon = if let Some(icon) = &agent_entry.icon {                            let icon_path = extensions_dir.join(ext_id).join(icon);                            // Canonicalize to resolve symlinks (dev extensions are symlinked)                            let absolute_icon_path = icon_path                                .canonicalize()                                .unwrap_or(icon_path)                                .to_string_lossy()                                .to_string();                             // Store icon locally for remote client                            self.agent_icons.insert(                                ExternalAgentServerName(agent_name.clone().into()),                                SharedString::from(absolute_icon_path.clone()),                            );                             Some(absolute_icon_path)                        } else {                            None                        };                         agents.push(ExternalExtensionAgent {                            name: agent_name.to_string(),                            icon_path: icon,                            extension_id: ext_id.to_string(),                            targets: agent_entry                                .targets                                .iter()                                .map(|(k, v)| (k.clone(), v.to_proto()))                                .collect(),                            env: agent_entry                                .env                                .iter()                                .map(|(k, v)| (k.clone(), v.clone()))                                .collect(),                        });                    }                }                upstream_client                    .read(cx)                    .proto_client()                    .send(proto::ExternalExtensionAgentsUpdated {                        project_id: *project_id,                        agents,                    })                    .log_err();            }            AgentServerStoreState::Collab => {                // Do nothing            }        }         cx.emit(AgentServersUpdated);    }     pub fn agent_icon(&self, name: &ExternalAgentServerName) -> Option<SharedString> {        self.agent_icons.get(name).cloned()    }     pub fn agent_display_name(&self, name: &ExternalAgentServerName) -> Option<SharedString> {        self.agent_display_names.get(name).cloned()    }     pub fn init_remote(session: &AnyProtoClient) {        session.add_entity_message_handler(Self::handle_external_agents_updated);        session.add_entity_message_handler(Self::handle_loading_status_updated);        session.add_entity_message_handler(Self::handle_new_version_available);    }     pub fn init_headless(session: &AnyProtoClient) {        session.add_entity_message_handler(Self::handle_external_extension_agents_updated);        session.add_entity_request_handler(Self::handle_get_agent_server_command);    }     fn agent_servers_settings_changed(&mut self, cx: &mut Context<Self>) {        let AgentServerStoreState::Local {            settings: old_settings,            ..        } = &mut self.state        else {            debug_panic!(                "should not be subscribed to agent server settings changes in non-local project"            );            return;        };         let new_settings = cx            .global::<SettingsStore>()            .get::<AllAgentServersSettings>(None)            .clone();        if Some(&new_settings) == old_settings.as_ref() {            return;        }         self.reregister_agents(cx);    }     fn reregister_agents(&mut self, cx: &mut Context<Self>) {        let AgentServerStoreState::Local {            node_runtime,            fs,            project_environment,            downstream_client,            settings: old_settings,            http_client,            extension_agents,            ..        } = &mut self.state        else {            debug_panic!("Non-local projects should never attempt to reregister. This is a bug!");             return;        };         let new_settings = cx            .global::<SettingsStore>()            .get::<AllAgentServersSettings>(None)            .clone();         self.external_agents.clear();        self.external_agents.insert(            GEMINI_NAME.into(),            Box::new(LocalGemini {                fs: fs.clone(),                node_runtime: node_runtime.clone(),                project_environment: project_environment.clone(),                custom_command: new_settings                    .gemini                    .clone()                    .and_then(|settings| settings.custom_command()),                ignore_system_version: new_settings                    .gemini                    .as_ref()                    .and_then(|settings| settings.ignore_system_version)                    .unwrap_or(true),            }),        );        self.external_agents.insert(            CODEX_NAME.into(),            Box::new(LocalCodex {                fs: fs.clone(),                project_environment: project_environment.clone(),                custom_command: new_settings                    .codex                    .clone()                    .and_then(|settings| settings.custom_command()),                http_client: http_client.clone(),                no_browser: downstream_client                    .as_ref()                    .is_some_and(|(_, client)| !client.has_wsl_interop()),            }),        );        self.external_agents.insert(            CLAUDE_CODE_NAME.into(),            Box::new(LocalClaudeCode {                fs: fs.clone(),                node_runtime: node_runtime.clone(),                project_environment: project_environment.clone(),                custom_command: new_settings                    .claude                    .clone()                    .and_then(|settings| settings.custom_command()),            }),        );        self.external_agents            .extend(                new_settings                    .custom                    .iter()                    .filter_map(|(name, settings)| match settings {                        CustomAgentServerSettings::Custom { command, .. } => Some((                            ExternalAgentServerName(name.clone()),

Lines 1–500 of 2443.