892 lines · 30.6 KB · Rust
RawDownload
use std::borrow::Cow;use std::cell::{Ref, RefCell};use std::path::{Path, PathBuf};use std::rc::Rc;use std::sync::Arc; use anyhow::{Context as _, Result, anyhow};use assistant_settings::{AgentProfile, AgentProfileId, AssistantSettings};use assistant_tool::{ToolId, ToolSource, ToolWorkingSet};use chrono::{DateTime, Utc};use collections::HashMap;use context_server::manager::ContextServerManager;use context_server::{ContextServerFactoryRegistry, ContextServerTool};use fs::Fs;use futures::channel::{mpsc, oneshot};use futures::future::{self, BoxFuture, Shared};use futures::{FutureExt as _, StreamExt as _};use gpui::{    App, BackgroundExecutor, Context, Entity, EventEmitter, Global, ReadGlobal, SharedString,    Subscription, Task, prelude::*,};use heed::Database;use heed::types::SerdeBincode;use language_model::{LanguageModelToolUseId, Role, TokenUsage};use project::{Project, Worktree};use prompt_store::{    ProjectContext, PromptBuilder, PromptId, PromptMetadata, PromptStore, PromptsUpdatedEvent,    RulesFileContext, UserPromptId, UserRulesContext, WorktreeContext,};use serde::{Deserialize, Serialize};use settings::{Settings as _, SettingsStore};use util::ResultExt as _; use crate::thread::{    DetailedSummaryState, ExceededWindowError, MessageId, ProjectSnapshot, Thread, ThreadId,}; const RULES_FILE_NAMES: [&'static str; 6] = [    ".rules",    ".cursorrules",    ".windsurfrules",    ".clinerules",    ".github/copilot-instructions.md",    "CLAUDE.md",]; pub fn init(cx: &mut App) {    ThreadsDatabase::init(cx);} /// A system prompt shared by all threads created by this ThreadStore#[derive(Clone, Default)]pub struct SharedProjectContext(Rc<RefCell<Option<ProjectContext>>>); impl SharedProjectContext {    pub fn borrow(&self) -> Ref<Option<ProjectContext>> {        self.0.borrow()    }} pub struct ThreadStore {    project: Entity<Project>,    tools: Entity<ToolWorkingSet>,    prompt_builder: Arc<PromptBuilder>,    prompt_store: Option<Entity<PromptStore>>,    context_server_manager: Entity<ContextServerManager>,    context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,    threads: Vec<SerializedThreadMetadata>,    project_context: SharedProjectContext,    reload_system_prompt_tx: mpsc::Sender<()>,    _reload_system_prompt_task: Task<()>,    _subscriptions: Vec<Subscription>,} pub struct RulesLoadingError {    pub message: SharedString,} impl EventEmitter<RulesLoadingError> for ThreadStore {} impl ThreadStore {    pub fn load(        project: Entity<Project>,        tools: Entity<ToolWorkingSet>,        prompt_builder: Arc<PromptBuilder>,        cx: &mut App,    ) -> Task<Result<Entity<Self>>> {        let prompt_store = PromptStore::global(cx);        cx.spawn(async move |cx| {            let prompt_store = prompt_store.await.ok();            let (thread_store, ready_rx) = cx.update(|cx| {                let mut option_ready_rx = None;                let thread_store = cx.new(|cx| {                    let (thread_store, ready_rx) =                        Self::new(project, tools, prompt_builder, prompt_store, cx);                    option_ready_rx = Some(ready_rx);                    thread_store                });                (thread_store, option_ready_rx.take().unwrap())            })?;            ready_rx.await?;            Ok(thread_store)        })    }     fn new(        project: Entity<Project>,        tools: Entity<ToolWorkingSet>,        prompt_builder: Arc<PromptBuilder>,        prompt_store: Option<Entity<PromptStore>>,        cx: &mut Context<Self>,    ) -> (Self, oneshot::Receiver<()>) {        let context_server_factory_registry = ContextServerFactoryRegistry::default_global(cx);        let context_server_manager = cx.new(|cx| {            ContextServerManager::new(context_server_factory_registry, project.clone(), cx)        });         let mut subscriptions = vec![            cx.observe_global::<SettingsStore>(move |this: &mut Self, cx| {                this.load_default_profile(cx);            }),            cx.subscribe(&project, Self::handle_project_event),        ];         if let Some(prompt_store) = prompt_store.as_ref() {            subscriptions.push(cx.subscribe(                prompt_store,                |this, _prompt_store, PromptsUpdatedEvent, _cx| {                    this.enqueue_system_prompt_reload();                },            ))        }         // This channel and task prevent concurrent and redundant loading of the system prompt.        let (reload_system_prompt_tx, mut reload_system_prompt_rx) = mpsc::channel(1);        let (ready_tx, ready_rx) = oneshot::channel();        let mut ready_tx = Some(ready_tx);        let reload_system_prompt_task = cx.spawn({            let prompt_store = prompt_store.clone();            async move |thread_store, cx| {                loop {                    let Some(reload_task) = thread_store                        .update(cx, |thread_store, cx| {                            thread_store.reload_system_prompt(prompt_store.clone(), cx)                        })                        .ok()                    else {                        return;                    };                    reload_task.await;                    if let Some(ready_tx) = ready_tx.take() {                        ready_tx.send(()).ok();                    }                    reload_system_prompt_rx.next().await;                }            }        });         let this = Self {            project,            tools,            prompt_builder,            prompt_store,            context_server_manager,            context_server_tool_ids: HashMap::default(),            threads: Vec::new(),            project_context: SharedProjectContext::default(),            reload_system_prompt_tx,            _reload_system_prompt_task: reload_system_prompt_task,            _subscriptions: subscriptions,        };        this.load_default_profile(cx);        this.register_context_server_handlers(cx);        this.reload(cx).detach_and_log_err(cx);        (this, ready_rx)    }     fn handle_project_event(        &mut self,        _project: Entity<Project>,        event: &project::Event,        _cx: &mut Context<Self>,    ) {        match event {            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {                self.enqueue_system_prompt_reload();            }            project::Event::WorktreeUpdatedEntries(_, items) => {                if items.iter().any(|(path, _, _)| {                    RULES_FILE_NAMES                        .iter()                        .any(|name| path.as_ref() == Path::new(name))                }) {                    self.enqueue_system_prompt_reload();                }            }            _ => {}        }    }     fn enqueue_system_prompt_reload(&mut self) {        self.reload_system_prompt_tx.try_send(()).ok();    }     // Note that this should only be called from `reload_system_prompt_task`.    fn reload_system_prompt(        &self,        prompt_store: Option<Entity<PromptStore>>,        cx: &mut Context<Self>,    ) -> Task<()> {        let project = self.project.read(cx);        let worktree_tasks = project            .visible_worktrees(cx)            .map(|worktree| {                Self::load_worktree_info_for_system_prompt(                    project.fs().clone(),                    worktree.read(cx),                    cx,                )            })            .collect::<Vec<_>>();        let default_user_rules_task = match prompt_store {            None => Task::ready(vec![]),            Some(prompt_store) => prompt_store.read_with(cx, |prompt_store, cx| {                let prompts = prompt_store.default_prompt_metadata();                let load_tasks = prompts.into_iter().map(|prompt_metadata| {                    let contents = prompt_store.load(prompt_metadata.id, cx);                    async move { (contents.await, prompt_metadata) }                });                cx.background_spawn(future::join_all(load_tasks))            }),        };         cx.spawn(async move |this, cx| {            let (worktrees, default_user_rules) =                future::join(future::join_all(worktree_tasks), default_user_rules_task).await;             let worktrees = worktrees                .into_iter()                .map(|(worktree, rules_error)| {                    if let Some(rules_error) = rules_error {                        this.update(cx, |_, cx| cx.emit(rules_error)).ok();                    }                    worktree                })                .collect::<Vec<_>>();             let default_user_rules = default_user_rules                .into_iter()                .flat_map(|(contents, prompt_metadata)| match contents {                    Ok(contents) => Some(UserRulesContext {                        uuid: match prompt_metadata.id {                            PromptId::User { uuid } => uuid,                            PromptId::EditWorkflow => return None,                        },                        title: prompt_metadata.title.map(|title| title.to_string()),                        contents,                    }),                    Err(err) => {                        this.update(cx, |_, cx| {                            cx.emit(RulesLoadingError {                                message: format!("{err:?}").into(),                            });                        })                        .ok();                        None                    }                })                .collect::<Vec<_>>();             this.update(cx, |this, _cx| {                *this.project_context.0.borrow_mut() =                    Some(ProjectContext::new(worktrees, default_user_rules));            })            .ok();        })    }     fn load_worktree_info_for_system_prompt(        fs: Arc<dyn Fs>,        worktree: &Worktree,        cx: &App,    ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {        let root_name = worktree.root_name().into();         let rules_task = Self::load_worktree_rules_file(fs, worktree, cx);        let Some(rules_task) = rules_task else {            return Task::ready((                WorktreeContext {                    root_name,                    rules_file: None,                },                None,            ));        };         cx.spawn(async move |_| {            let (rules_file, rules_file_error) = match rules_task.await {                Ok(rules_file) => (Some(rules_file), None),                Err(err) => (                    None,                    Some(RulesLoadingError {                        message: format!("{err}").into(),                    }),                ),            };            let worktree_info = WorktreeContext {                root_name,                rules_file,            };            (worktree_info, rules_file_error)        })    }     fn load_worktree_rules_file(        fs: Arc<dyn Fs>,        worktree: &Worktree,        cx: &App,    ) -> Option<Task<Result<RulesFileContext>>> {        let selected_rules_file = RULES_FILE_NAMES            .into_iter()            .filter_map(|name| {                worktree                    .entry_for_path(name)                    .filter(|entry| entry.is_file())                    .map(|entry| (entry.path.clone(), worktree.absolutize(&entry.path)))            })            .next();         // Note that Cline supports `.clinerules` being a directory, but that is not currently        // supported. This doesn't seem to occur often in GitHub repositories.        selected_rules_file.map(|(path_in_worktree, abs_path)| {            let fs = fs.clone();            cx.background_spawn(async move {                let abs_path = abs_path?;                let text = fs.load(&abs_path).await.with_context(|| {                    format!("Failed to load assistant rules file {:?}", abs_path)                })?;                anyhow::Ok(RulesFileContext {                    path_in_worktree,                    abs_path: abs_path.into(),                    text: text.trim().to_string(),                })            })        })    }     pub fn context_server_manager(&self) -> Entity<ContextServerManager> {        self.context_server_manager.clone()    }     pub fn prompt_store(&self) -> Option<Entity<PromptStore>> {        self.prompt_store.clone()    }     pub fn load_rules(        &self,        prompt_id: UserPromptId,        cx: &App,    ) -> Task<Result<(PromptMetadata, String)>> {        let prompt_id = PromptId::User { uuid: prompt_id };        let Some(prompt_store) = self.prompt_store.as_ref() else {            return Task::ready(Err(anyhow!("Prompt store unexpectedly missing.")));        };        let prompt_store = prompt_store.read(cx);        let Some(metadata) = prompt_store.metadata(prompt_id) else {            return Task::ready(Err(anyhow!("User rules not found in library.")));        };        let text_task = prompt_store.load(prompt_id, cx);        cx.background_spawn(async move { Ok((metadata, text_task.await?)) })    }     pub fn tools(&self) -> Entity<ToolWorkingSet> {        self.tools.clone()    }     /// Returns the number of threads.    pub fn thread_count(&self) -> usize {        self.threads.len()    }     pub fn threads(&self) -> Vec<SerializedThreadMetadata> {        let mut threads = self.threads.iter().cloned().collect::<Vec<_>>();        threads.sort_unstable_by_key(|thread| std::cmp::Reverse(thread.updated_at));        threads    }     pub fn recent_threads(&self, limit: usize) -> Vec<SerializedThreadMetadata> {        self.threads().into_iter().take(limit).collect()    }     pub fn create_thread(&mut self, cx: &mut Context<Self>) -> Entity<Thread> {        cx.new(|cx| {            Thread::new(                self.project.clone(),                self.tools.clone(),                self.prompt_builder.clone(),                self.project_context.clone(),                cx,            )        })    }     pub fn open_thread(        &self,        id: &ThreadId,        cx: &mut Context<Self>,    ) -> Task<Result<Entity<Thread>>> {        let id = id.clone();        let database_future = ThreadsDatabase::global_future(cx);        cx.spawn(async move |this, cx| {            let database = database_future.await.map_err(|err| anyhow!(err))?;            let thread = database                .try_find_thread(id.clone())                .await?                .ok_or_else(|| anyhow!("no thread found with ID: {id:?}"))?;             let thread = this.update(cx, |this, cx| {                cx.new(|cx| {                    Thread::deserialize(                        id.clone(),                        thread,                        this.project.clone(),                        this.tools.clone(),                        this.prompt_builder.clone(),                        this.project_context.clone(),                        cx,                    )                })            })?;             Ok(thread)        })    }     pub fn save_thread(&self, thread: &Entity<Thread>, cx: &mut Context<Self>) -> Task<Result<()>> {        let (metadata, serialized_thread) =            thread.update(cx, |thread, cx| (thread.id().clone(), thread.serialize(cx)));         let database_future = ThreadsDatabase::global_future(cx);        cx.spawn(async move |this, cx| {            let serialized_thread = serialized_thread.await?;            let database = database_future.await.map_err(|err| anyhow!(err))?;            database.save_thread(metadata, serialized_thread).await?;             this.update(cx, |this, cx| this.reload(cx))?.await        })    }     pub fn delete_thread(&mut self, id: &ThreadId, cx: &mut Context<Self>) -> Task<Result<()>> {        let id = id.clone();        let database_future = ThreadsDatabase::global_future(cx);        cx.spawn(async move |this, cx| {            let database = database_future.await.map_err(|err| anyhow!(err))?;            database.delete_thread(id.clone()).await?;             this.update(cx, |this, cx| {                this.threads.retain(|thread| thread.id != id);                cx.notify();            })        })    }     pub fn reload(&self, cx: &mut Context<Self>) -> Task<Result<()>> {        let database_future = ThreadsDatabase::global_future(cx);        cx.spawn(async move |this, cx| {            let threads = database_future                .await                .map_err(|err| anyhow!(err))?                .list_threads()                .await?;             this.update(cx, |this, cx| {                this.threads = threads;                cx.notify();            })        })    }     fn load_default_profile(&self, cx: &mut Context<Self>) {        let assistant_settings = AssistantSettings::get_global(cx);         self.load_profile_by_id(assistant_settings.default_profile.clone(), cx);    }     pub fn load_profile_by_id(&self, profile_id: AgentProfileId, cx: &mut Context<Self>) {        let assistant_settings = AssistantSettings::get_global(cx);         if let Some(profile) = assistant_settings.profiles.get(&profile_id) {            self.load_profile(profile.clone(), cx);        }    }     pub fn load_profile(&self, profile: AgentProfile, cx: &mut Context<Self>) {        self.tools.update(cx, |tools, cx| {            tools.disable_all_tools(cx);            tools.enable(                ToolSource::Native,                &profile                    .tools

Lines 1–500 of 892.