5331 lines · 200.3 KB · Rust
RawDownload
mod db;mod legacy_thread;mod native_agent_server;pub mod outline;mod pattern_extraction;mod sandboxing;mod templates;#[cfg(test)]mod tests;mod thread;mod thread_store;mod tool_permissions;mod tools; use context_server::ContextServerId;pub use db::*;use itertools::Itertools;pub use native_agent_server::NativeAgentServer;pub use pattern_extraction::*;pub use shell_command_parser::extract_commands;pub use templates::*;pub use thread::*;pub use thread_store::*;pub use tool_permissions::*;pub use tools::*; use acp_thread::{    AcpThread, AgentModelSelector, AgentSessionInfo, AgentSessionList, AgentSessionListRequest,    AgentSessionListResponse, TokenUsageRatio, UserMessageId,};use agent_client_protocol::schema as acp;use agent_skills::{    MAX_SKILL_DESCRIPTIONS_SIZE, ProjectSkillGroup, Skill, SkillIndex, SkillLoadError,    SkillScopeId, SkillSource, SkillSummary, builtin_skills, global_skills_dir,    load_skills_from_directory, project_skills_relative_path,};use anyhow::{Context as _, Result, anyhow};use chrono::{DateTime, Utc};use collections::{HashMap, HashSet, IndexMap}; use fs::Fs;use futures::channel::{mpsc, oneshot};use futures::future::Shared;use futures::{FutureExt as _, StreamExt as _, future};use gpui::{    App, AppContext, AsyncApp, Context, Entity, EntityId, SharedString, Subscription, Task,    TaskExt, WeakEntity,};use language_model::{IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelRegistry};use project::{    AgentId, Project, ProjectItem, ProjectPath, Worktree, trusted_worktrees::TrustedWorktrees,};use prompt_store::{ProjectContext, RULES_FILE_NAMES, RulesFileContext, WorktreeContext};use serde::{Deserialize, Serialize};use settings::{LanguageModelSelection, Settings as _, update_settings_file};use std::any::Any;use std::path::PathBuf;use std::rc::Rc;use std::sync::{Arc, LazyLock};use util::ResultExt;use util::path_list::PathList;use util::rel_path::RelPath; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]pub struct ProjectSnapshot {    pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,    pub timestamp: DateTime<Utc>,} pub struct RulesLoadingError {    pub message: SharedString,} #[derive(Clone, Debug, PartialEq, Eq, Hash)]pub struct SkillLoadingError {    pub project_id: EntityId,    pub path: PathBuf,    pub message: SharedString,} /// Emitted whenever the set of skill loading errors for a project changes./// The `errors` field is the full replacement list; subscribers should treat/// it as a snapshot rather than appending. An empty `errors` list means all/// previously-reported errors have been resolved.#[derive(Clone, Debug)]pub struct SkillLoadingErrorsUpdated {    pub project_id: EntityId,    pub errors: Vec<SkillLoadingError>,} #[derive(Clone, Debug)]pub struct NativeAvailableSkill {    pub name: String,    pub description: String,    pub source: SharedString,    pub skill_file_path: PathBuf,} impl From<&Skill> for NativeAvailableSkill {    fn from(skill: &Skill) -> Self {        Self {            name: skill.name.clone(),            description: skill.description.clone(),            source: skill.source.display_label().to_string().into(),            skill_file_path: skill.skill_file_path.clone(),        }    }} struct ProjectState {    project: Entity<Project>,    project_context: Entity<ProjectContext>,    skills: Arc<Vec<Skill>>,    skill_loading_errors: Vec<SkillLoadingError>,    project_context_needs_refresh: watch::Sender<()>,    _maintain_project_context: Task<Result<()>>,    context_server_registry: Entity<ContextServerRegistry>,    _subscriptions: Vec<Subscription>,} /// Holds both the internal Thread and the AcpThread for a sessionstruct Session {    /// The internal thread that processes messages    thread: Entity<Thread>,    /// The ACP thread that handles protocol communication    acp_thread: Entity<acp_thread::AcpThread>,    project_id: EntityId,    pending_save: Task<Result<()>>,    _subscriptions: Vec<Subscription>,    ref_count: usize,} struct PendingSession {    task: Shared<Task<Result<Entity<AcpThread>, Arc<anyhow::Error>>>>,    ref_count: usize,} pub struct LanguageModels {    /// Access language model by ID    models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,    /// Cached list for returning language model information    model_list: acp_thread::AgentModelList,    refresh_models_rx: watch::Receiver<()>,    refresh_models_tx: watch::Sender<()>,    _authenticate_all_providers_task: Task<()>,} impl LanguageModels {    fn new(cx: &mut App) -> Self {        let (refresh_models_tx, refresh_models_rx) = watch::channel(());         let mut this = Self {            models: HashMap::default(),            model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),            refresh_models_rx,            refresh_models_tx,            _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),        };        this.refresh_list(cx);        this    }     fn refresh_list(&mut self, cx: &App) {        let providers = LanguageModelRegistry::global(cx)            .read(cx)            .visible_providers()            .into_iter()            .filter(|provider| provider.is_authenticated(cx))            .collect::<Vec<_>>();         let mut language_model_list = IndexMap::default();        let mut recommended_models = HashSet::default();         let mut recommended = Vec::new();        for provider in &providers {            for model in provider.recommended_models(cx) {                recommended_models.insert((model.provider_id(), model.id()));                recommended.push(Self::map_language_model_to_info(&model, provider));            }        }        if !recommended.is_empty() {            language_model_list.insert(                acp_thread::AgentModelGroupName("Recommended".into()),                recommended,            );        }         let mut models = HashMap::default();        for provider in providers {            let mut provider_models = Vec::new();            for model in provider.provided_models(cx) {                let model_info = Self::map_language_model_to_info(&model, &provider);                let model_id = model_info.id.clone();                provider_models.push(model_info);                models.insert(model_id, model);            }            if !provider_models.is_empty() {                language_model_list.insert(                    acp_thread::AgentModelGroupName(provider.name().0.clone()),                    provider_models,                );            }        }         self.models = models;        self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);        self.refresh_models_tx.send(()).ok();    }     fn watch(&self) -> watch::Receiver<()> {        self.refresh_models_rx.clone()    }     pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {        self.models.get(model_id).cloned()    }     fn map_language_model_to_info(        model: &Arc<dyn LanguageModel>,        provider: &Arc<dyn LanguageModelProvider>,    ) -> acp_thread::AgentModelInfo {        acp_thread::AgentModelInfo {            id: Self::model_id(model),            name: model.name().0,            description: None,            icon: Some(match provider.icon() {                IconOrSvg::Svg(path) => acp_thread::AgentModelIcon::Path(path),                IconOrSvg::Icon(name) => acp_thread::AgentModelIcon::Named(name),            }),            is_latest: model.is_latest(),            cost: model.model_cost_info().map(|cost| cost.to_shared_string()),        }    }     fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {        acp::ModelId::new(format!("{}/{}", model.provider_id().0, model.id().0))    }     fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {        let authenticate_all_providers = LanguageModelRegistry::global(cx)            .read(cx)            .visible_providers()            .iter()            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))            .collect::<Vec<_>>();         cx.spawn(async move |cx| {            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {                if let Err(err) = authenticate_task.await {                    match err {                        language_model::AuthenticateError::CredentialsNotFound => {                            // Since we're authenticating these providers in the                            // background for the purposes of populating the                            // language selector, we don't care about providers                            // where the credentials are not found.                        }                        language_model::AuthenticateError::ConnectionRefused => {                            // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.                            // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.                            // TODO: Better manage LM Studio auth logic to avoid these noisy failures.                        }                        _ => {                            // Some providers have noisy failure states that we                            // don't want to spam the logs with every time the                            // language model selector is initialized.                            //                            // Ideally these should have more clear failure modes                            // that we know are safe to ignore here, like what we do                            // with `CredentialsNotFound` above.                            match provider_id.0.as_ref() {                                "lmstudio" | "ollama" => {                                    // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".                                    //                                    // These fail noisily, so we don't log them.                                }                                "copilot_chat" => {                                    // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.                                }                                _ => {                                    log::error!(                                        "Failed to authenticate provider: {}: {err:#}",                                        provider_name.0                                    );                                }                            }                        }                    }                }            }             cx.update(language_models::update_environment_fallback_model);        })    }} pub struct NativeAgent {    /// Session ID -> Session mapping    sessions: HashMap<acp::SessionId, Session>,    pending_sessions: HashMap<acp::SessionId, PendingSession>,    thread_store: Entity<ThreadStore>,    /// Project-specific state keyed by project EntityId    projects: HashMap<EntityId, ProjectState>,    /// Shared templates for all threads    templates: Arc<Templates>,    /// Cached model information    models: LanguageModels,    fs: Arc<dyn Fs>,    _subscriptions: Vec<Subscription>,    /// Tracks the lifecycle of global skills directory observation. We    /// don't eagerly watch (or even check for) `~/.agents/skills/` at    /// startup; users who never engage with the agent panel pay zero    /// filesystem cost. The watch is kicked off lazily by    /// [`Self::ensure_skills_scan_started`], which is called from the    /// three agent-panel interaction points: input box focus, slash    /// autocomplete, and conversation submit.    skills_state: SkillsState,} #[derive(Default)]enum SkillsState {    /// No scan or watch is active. A user-interaction trigger will kick    /// off a fresh scan.    #[default]    Idle,    /// A one-shot scan task is in flight. It checks whether    /// `~/.agents/skills/` exists; if so, transitions to `Watching`,    /// otherwise back to `Idle`.    Scanning,    /// A watch task is observing `~/.agents/skills/`. It transitions    /// back to `Idle` if the watched directory itself is removed.    Watching,} impl gpui::EventEmitter<SkillLoadingErrorsUpdated> for NativeAgent {} static RULES_FILE_REL_PATHS: LazyLock<Vec<Arc<RelPath>>> = LazyLock::new(|| {    RULES_FILE_NAMES        .iter()        .filter_map(|name| RelPath::unix(name).ok().map(|path| path.into_arc()))        .collect()}); static SKILLS_PREFIX: LazyLock<Option<Arc<RelPath>>> = LazyLock::new(|| {    RelPath::unix(project_skills_relative_path())        .ok()        .map(|path| path.into_arc())}); impl NativeAgent {    pub fn new(        thread_store: Entity<ThreadStore>,        templates: Arc<Templates>,        fs: Arc<dyn Fs>,        cx: &mut App,    ) -> Entity<NativeAgent> {        log::debug!("Creating new NativeAgent");         cx.new(|cx| {            let subscriptions = vec![cx.subscribe(                &LanguageModelRegistry::global(cx),                Self::handle_models_updated_event,            )];             if !cx.has_global::<SkillIndex>() {                cx.set_global(SkillIndex::default());            }             Self {                sessions: HashMap::default(),                pending_sessions: HashMap::default(),                thread_store,                projects: HashMap::default(),                templates,                models: LanguageModels::new(cx),                fs,                _subscriptions: subscriptions,                skills_state: SkillsState::default(),            }        })    }     /// Kicks off a one-time scan of the global skills directory if one    /// isn't already in progress and a watch isn't already active.    ///    /// Idempotent and cheap: returns immediately if a scan or watch is    /// already running. The expected callers are user-interaction events    /// from the agent panel (input focus, slash autocomplete, conversation    /// submit); firing this from any of them is equivalent and safe to    /// repeat.    ///    /// The scan itself runs detached on the foreground executor. If    /// `~/.agents/skills/` exists it transitions state to    /// [`SkillsState::Watching`] and starts a recursive watch;    /// otherwise it transitions back to [`SkillsState::Idle`] so the    /// next trigger retries (covering the case where the user creates    /// the directory after the first scan).    pub fn ensure_skills_scan_started(&mut self, cx: &mut Context<Self>) {        if !matches!(self.skills_state, SkillsState::Idle) {            return;        }        self.skills_state = SkillsState::Scanning;        let fs = self.fs.clone();        cx.spawn(async move |this, cx| Self::run_skills_scan(this, fs, cx).await)            .detach();    }     async fn run_skills_scan(this: WeakEntity<Self>, fs: Arc<dyn Fs>, cx: &mut AsyncApp) {        let skills_dir = global_skills_dir();        if !fs.is_dir(&skills_dir).await {            // Skills directory doesn't exist; revert state so the next            // user trigger retries.            let _ = this.update(cx, |this, _cx| {                this.skills_state = SkillsState::Idle;            });            return;        }         // Skills directory exists. Start a watch and trigger a refresh        // of every project's context so the freshly-discovered skills        // get loaded.        let _ = this.update(cx, |this, cx| {            cx.spawn({                let fs = fs.clone();                let skills_dir = skills_dir.clone();                async move |this, cx| Self::run_skills_watch(this, fs, skills_dir, cx).await            })            .detach();            this.skills_state = SkillsState::Watching;            for state in this.projects.values_mut() {                state.project_context_needs_refresh.send(()).ok();            }        });    }     async fn run_skills_watch(        this: WeakEntity<Self>,        fs: Arc<dyn Fs>,        skills_dir: PathBuf,        cx: &mut AsyncApp,    ) {        let (mut events, watcher) = fs            .watch(&skills_dir, std::time::Duration::from_millis(500))            .await;         // Linux's inotify backend is non-recursive, so a watch on        // `skills_dir` only fires for direct children. Skill discovery        // is intentionally one level deep (`<skills_dir>/<skill>/SKILL.md`),        // so we only register watches on each immediate child directory        // and deliberately do NOT recurse: a stray `node_modules`,        // `target`, or `.git` inside a skill folder would otherwise        // register watches for tens of thousands of subdirectories.        // These per-child adds are cheap no-ops on macOS/Windows where        // the OS-level watch is already recursive.        if let Ok(mut entries) = fs.read_dir(&skills_dir).await {            while let Some(entry) = entries.next().await {                let Ok(path) = entry else { continue };                if let Ok(Some(metadata)) = fs.metadata(&path).await                    && metadata.is_dir                {                    watcher.add(&path).ok();                }            }        }         while let Some(events) = events.next().await {            // When a new immediate child directory of `skills_dir` is            // created, add a single watch for it so changes to its            // `SKILL.md` are observed on Linux. We intentionally do not            // recurse into the new directory — skill discovery is only            // one level deep.            for event in &events {                if event.kind == Some(fs::PathEventKind::Created)                    && event.path.parent() == Some(skills_dir.as_path())                    && fs.is_dir(&event.path).await                {                    watcher.add(&event.path).ok();                }            }             let watched_root_removed = events.iter().any(|event| {                event.path == skills_dir && event.kind == Some(fs::PathEventKind::Removed)            });             let updated = this.update(cx, |this, _cx| {                for state in this.projects.values_mut() {                    state.project_context_needs_refresh.send(()).ok();                }                if watched_root_removed {                    // Drop back to Idle so the next user trigger                    // retries the scan; the next trigger will rediscover                    // the directory if the user has recreated it.                    this.skills_state = SkillsState::Idle;                }            });            if updated.is_err() || watched_root_removed {                return;            }        }    } 

Lines 1–500 of 5331.