4273 lines · 153.5 KB · Rust
RawDownload
use std::{    path::{Path, PathBuf},    sync::Arc,}; use agent::{ThreadStore, ZED_AGENT_ID};use agent_client_protocol::schema::v1 as acp;use anyhow::Context as _;use chrono::{DateTime, Utc};use collections::{HashMap, HashSet};use db::{    kvp::KeyValueStore,    sqlez::{        bindable::{Bind, Column},        domain::Domain,        statement::Statement,        thread_safe_connection::ThreadSafeConnection,    },    sqlez_macros::sql,};use fs::Fs;use futures::{FutureExt, future::Shared};use gpui::{AppContext as _, Entity, Global, Subscription, Task, TaskExt};pub use project::WorktreePaths;use project::{AgentId, linked_worktree_short_name};use remote::{RemoteConnectionOptions, same_remote_connection_identity};use ui::{App, Context, SharedString, ThreadItemWorktreeInfo, WorktreeKind};use util::ResultExt as _;use workspace::{PathList, SerializedWorkspaceLocation, WorkspaceDb}; use crate::DEFAULT_THREAD_TITLE; #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, serde::Serialize, serde::Deserialize)]pub struct ThreadId(uuid::Uuid); impl ThreadId {    pub fn new() -> Self {        Self(uuid::Uuid::new_v4())    }     /// Stable, hyphenated string form suitable for use as a key.    pub fn to_key_string(&self) -> String {        self.0.hyphenated().to_string()    }} impl Bind for ThreadId {    fn bind(&self, statement: &Statement, start_index: i32) -> anyhow::Result<i32> {        self.0.bind(statement, start_index)    }} impl Column for ThreadId {    fn column(statement: &mut Statement, start_index: i32) -> anyhow::Result<(Self, i32)> {        let (uuid, next) = Column::column(statement, start_index)?;        Ok((ThreadId(uuid), next))    }} const THREAD_REMOTE_CONNECTION_MIGRATION_KEY: &str = "thread-metadata-remote-connection-backfill";const THREAD_ID_MIGRATION_KEY: &str = "thread-metadata-thread-id-backfill"; /// List all sidebar thread metadata from an arbitrary SQLite connection.////// This is used to read thread metadata from another release channel's/// database without opening a full `ThreadSafeConnection`.pub(crate) fn list_thread_metadata_from_connection(    connection: &db::sqlez::connection::Connection,) -> anyhow::Result<Vec<ThreadMetadata>> {    connection.select::<ThreadMetadata>(ThreadMetadataDb::LIST_QUERY)?()} /// Run the `ThreadMetadataDb` migrations on a raw connection.////// This is used in tests to set up the sidebar_threads schema in a/// temporary database.#[cfg(test)]pub(crate) fn run_thread_metadata_migrations(connection: &db::sqlez::connection::Connection) {    connection        .migrate(            ThreadMetadataDb::NAME,            ThreadMetadataDb::MIGRATIONS,            &mut |_, _, _| false,        )        .expect("thread metadata migrations should succeed");} pub fn init(cx: &mut App) {    ThreadMetadataStore::init_global(cx);    let migration_task = migrate_thread_metadata(cx);    migrate_thread_remote_connections(cx, migration_task);    migrate_thread_ids(cx);} /// Migrate existing thread metadata from native agent thread store to the new metadata storage./// We skip migrating threads that do not have a project.////// TODO: Remove this after N weeks of shipping the sidebarfn migrate_thread_metadata(cx: &mut App) -> Task<anyhow::Result<()>> {    let store = ThreadMetadataStore::global(cx);    let db = store.read(cx).db.clone();    let thread_store = ThreadStore::global(cx);    let thread_store_ready = thread_store.read(cx).reload_task();     cx.spawn(async move |cx| {        // Wait for `ThreadStore`'s initial reload to complete. Without this,        // reading `entries()` races with the store's async population from        // disk and usually observes an empty iterator, silently skipping the        // migration on every launch. The regression test        // `test_migration_awaits_thread_store_reload` pins this behavior.        thread_store_ready.await;         let existing_list = db.list()?;        let existing_session_ids: HashSet<Arc<str>> = existing_list            .into_iter()            .filter_map(|m| m.session_id.map(|s| s.0))            .collect();         let mut to_migrate = thread_store.read_with(cx, |store, _cx| {            store                .entries()                .filter_map(|entry| {                    if existing_session_ids.contains(&entry.id.0) {                        return None;                    }                     Some(ThreadMetadata {                        thread_id: ThreadId::new(),                        session_id: Some(entry.id),                        agent_id: ZED_AGENT_ID.clone(),                        title: if entry.title.is_empty()                            || entry.title.as_ref() == DEFAULT_THREAD_TITLE                        {                            None                        } else {                            Some(entry.title)                        },                        title_override: None,                        updated_at: entry.updated_at,                        created_at: entry.created_at,                        interacted_at: None,                        worktree_paths: WorktreePaths::from_folder_paths(&entry.folder_paths),                        remote_connection: None,                        archived: true,                    })                })                .collect::<Vec<_>>()        });         if to_migrate.is_empty() {            return anyhow::Ok(());        }         // For each batch of newly-migrated threads, keep the 5 most recent        // per project unarchived. Previously this was gated on        // `is_first_migration` (an empty `sidebar_threads`), which meant any        // subsequent batch of newly-discovered legacy threads got migrated as        // fully archived. Running the rescue per-batch keeps the behavior        // idempotent across partial migrations and re-runs.        let mut per_project: HashMap<PathList, Vec<&mut ThreadMetadata>> = HashMap::default();        for entry in &mut to_migrate {            if entry.worktree_paths.is_empty() {                continue;            }            per_project                .entry(entry.worktree_paths.folder_path_list().clone())                .or_default()                .push(entry);        }        for entries in per_project.values_mut() {            entries.sort_by_key(|entry| std::cmp::Reverse(entry.updated_at));            for entry in entries.iter_mut().take(5) {                entry.archived = false;            }        }         log::info!("Migrating {} thread store entries", to_migrate.len());         // Manually save each entry to the database and call reload, otherwise        // we'll end up triggering lots of reloads after each save        for entry in to_migrate {            db.save(entry).await?;        }         log::info!("Finished migrating thread store entries");         let _ = store.update(cx, |store, cx| store.reload(cx));        anyhow::Ok(())    })} fn migrate_thread_remote_connections(cx: &mut App, migration_task: Task<anyhow::Result<()>>) {    let store = ThreadMetadataStore::global(cx);    let db = store.read(cx).db.clone();    let kvp = KeyValueStore::global(cx);    let workspace_db = WorkspaceDb::global(cx);    let fs = <dyn Fs>::global(cx);     cx.spawn(async move |cx| -> anyhow::Result<()> {        migration_task.await?;         if kvp            .read_kvp(THREAD_REMOTE_CONNECTION_MIGRATION_KEY)?            .is_some()        {            return Ok(());        }         let recent_workspaces = workspace_db            .recent_project_workspaces_ungrouped(fs.as_ref())            .await?;         let mut local_path_lists = HashSet::<PathList>::default();        let mut remote_path_lists = HashMap::<PathList, RemoteConnectionOptions>::default();         recent_workspaces            .iter()            .filter(|workspace| {                !workspace.paths.is_empty()                    && matches!(workspace.location, SerializedWorkspaceLocation::Local)            })            .for_each(|workspace| {                local_path_lists.insert(workspace.paths.clone());            });         for workspace in recent_workspaces {            match workspace.location {                SerializedWorkspaceLocation::Remote(remote_connection)                    if !local_path_lists.contains(&workspace.paths) =>                {                    remote_path_lists                        .entry(workspace.paths)                        .or_insert(remote_connection);                }                _ => {}            }        }         let mut reloaded = false;        for metadata in db.list()? {            if metadata.remote_connection.is_some() {                continue;            }             if let Some(remote_connection) = remote_path_lists                .get(metadata.folder_paths())                .or_else(|| remote_path_lists.get(metadata.main_worktree_paths()))            {                db.save(ThreadMetadata {                    remote_connection: Some(remote_connection.clone()),                    ..metadata                })                .await?;                reloaded = true;            }        }         let reloaded_task = reloaded            .then_some(store.update(cx, |store, cx| store.reload(cx)))            .unwrap_or(Task::ready(()).shared());         kvp.write_kvp(            THREAD_REMOTE_CONNECTION_MIGRATION_KEY.to_string(),            "1".to_string(),        )        .await?;        reloaded_task.await;         Ok(())    })    .detach_and_log_err(cx);} fn migrate_thread_ids(cx: &mut App) {    let store = ThreadMetadataStore::global(cx);    let db = store.read(cx).db.clone();    let kvp = KeyValueStore::global(cx);     cx.spawn(async move |cx| -> anyhow::Result<()> {        if kvp.read_kvp(THREAD_ID_MIGRATION_KEY)?.is_some() {            return Ok(());        }         let mut reloaded = false;        for metadata in db.list()? {            db.save(metadata).await?;            reloaded = true;        }         let reloaded_task = reloaded            .then_some(store.update(cx, |store, cx| store.reload(cx)))            .unwrap_or(Task::ready(()).shared());         kvp.write_kvp(THREAD_ID_MIGRATION_KEY.to_string(), "1".to_string())            .await?;        reloaded_task.await;         Ok(())    })    .detach_and_log_err(cx);} struct GlobalThreadMetadataStore(Entity<ThreadMetadataStore>);impl Global for GlobalThreadMetadataStore {} /// Lightweight metadata for any thread (native or ACP), enough to populate/// the sidebar list and route to the correct load path when clicked.#[derive(Debug, Clone, PartialEq)]pub struct ThreadMetadata {    pub thread_id: ThreadId,    pub session_id: Option<acp::SessionId>,    pub agent_id: AgentId,    pub title: Option<SharedString>,    /// User-supplied title that takes precedence over `title`. Set when the    /// user renames a thread, so that subsequent agent-driven title updates    /// (e.g. from `SessionInfoUpdate`) don't clobber the user's choice.    pub title_override: Option<SharedString>,    pub updated_at: DateTime<Utc>,    pub created_at: Option<DateTime<Utc>>,    /// When a user last interacted to send a message (including queueing).    /// Doesn't include the time when a queued message is fired.    pub interacted_at: Option<DateTime<Utc>>,    pub worktree_paths: WorktreePaths,    pub remote_connection: Option<RemoteConnectionOptions>,    pub archived: bool,} impl ThreadMetadata {    /// A thread is a draft until its first message is sent, at which point    /// it gets an ACP `session_id`.    pub fn is_draft(&self) -> bool {        self.session_id.is_none()    }     pub fn display_title(&self) -> SharedString {        self.title()            .unwrap_or_else(|| crate::DEFAULT_THREAD_TITLE.into())    }     pub fn title(&self) -> Option<SharedString> {        self.title_override.clone().or_else(|| self.title.clone())    }     pub fn folder_paths(&self) -> &PathList {        self.worktree_paths.folder_path_list()    }    pub fn main_worktree_paths(&self) -> &PathList {        self.worktree_paths.main_worktree_path_list()    }     pub fn references_folder_path(&self, path: &Path) -> bool {        self.folder_paths()            .paths()            .iter()            .any(|folder_path| folder_path.as_path() == path)    }     pub fn matches_remote_connection(        &self,        remote_connection: Option<&RemoteConnectionOptions>,    ) -> bool {        same_remote_connection_identity(self.remote_connection.as_ref(), remote_connection)    }} /// Derives worktree display info from a thread's stored path list.////// For each path in the thread's `folder_paths`, produces a/// [`ThreadItemWorktreeInfo`] with a short display name, full path, and whether/// the worktree is the main checkout or a linked git worktree. When/// multiple main paths exist and a linked worktree's short name alone/// wouldn't identify which main project it belongs to, the main project/// name is prefixed for disambiguation (e.g. `project:feature`).pub fn worktree_info_from_thread_paths<S: std::hash::BuildHasher>(    worktree_paths: &WorktreePaths,    branch_names: &std::collections::HashMap<PathBuf, SharedString, S>,) -> Vec<ThreadItemWorktreeInfo> {    let mut infos: Vec<ThreadItemWorktreeInfo> = Vec::new();    let mut linked_short_names: Vec<(SharedString, SharedString)> = Vec::new();    let mut unique_main_count = HashSet::default();     for (main_path, folder_path) in worktree_paths.ordered_pairs() {        unique_main_count.insert(main_path.clone());        let is_linked = main_path != folder_path;         if is_linked {            let short_name = linked_worktree_short_name(main_path, folder_path).unwrap_or_default();            let project_name = main_path                .file_name()                .map(|n| SharedString::from(n.to_string_lossy().to_string()))                .unwrap_or_default();            linked_short_names.push((short_name.clone(), project_name));            infos.push(ThreadItemWorktreeInfo {                worktree_name: Some(short_name),                full_path: SharedString::from(folder_path.display().to_string()),                highlight_positions: Vec::new(),                kind: WorktreeKind::Linked,                branch_name: branch_names.get(folder_path).cloned(),            });        } else {            let Some(name) = folder_path.file_name() else {                continue;            };            infos.push(ThreadItemWorktreeInfo {                worktree_name: Some(SharedString::from(name.to_string_lossy().to_string())),                full_path: SharedString::from(folder_path.display().to_string()),                highlight_positions: Vec::new(),                kind: WorktreeKind::Main,                branch_name: branch_names.get(folder_path).cloned(),            });        }    }     // When the group has multiple main worktree paths and the thread's    // folder paths don't all share the same short name, prefix each    // linked worktree chip with its main project name so the user knows    // which project it belongs to.    let all_same_name = infos.len() > 1        && infos            .iter()            .all(|i| i.worktree_name == infos[0].worktree_name);     if unique_main_count.len() > 1 && !all_same_name {        for (info, (_short_name, project_name)) in infos            .iter_mut()            .filter(|i| i.kind == WorktreeKind::Linked)            .zip(linked_short_names.iter())        {            if let Some(name) = &info.worktree_name {                info.worktree_name = Some(SharedString::from(format!("{}:{}", project_name, name)));            }        }    }     infos} impl From<&ThreadMetadata> for acp_thread::AgentSessionInfo {    fn from(meta: &ThreadMetadata) -> Self {        let session_id = meta            .session_id            .clone()            .unwrap_or_else(|| acp::SessionId::new(meta.thread_id.0.to_string()));        Self {            session_id,            work_dirs: Some(meta.folder_paths().clone()),            title: meta.title(),            updated_at: Some(meta.updated_at),            created_at: meta.created_at,            meta: None,        }    }} /// Record of a git worktree that was archived (deleted from disk) when its/// last thread was archived.pub struct ArchivedGitWorktree {    /// Auto-incrementing primary key.    pub id: i64,    /// Absolute path to the directory of the worktree before it was deleted.    /// Used when restoring, to put the recreated worktree back where it was.    /// If the path already exists on disk, the worktree is assumed to be    /// already restored and is used as-is.    pub worktree_path: PathBuf,    /// Absolute path of the main repository ("main worktree") that owned this worktree.    /// Used when restoring, to reattach the recreated worktree to the correct main repo.    /// If the main repo isn't found on disk, unarchiving fails because we only store    /// commit hashes, and without the actual git repo being available, we can't restore    /// the files.    pub main_repo_path: PathBuf,    /// Branch that was checked out in the worktree at archive time. `None` if    /// the worktree was in detached HEAD state, which isn't supported in Zed, but    /// could happen if the user made a detached one outside of Zed.    /// On restore, we try to switch to this branch. If that fails (e.g. it's    /// checked out elsewhere), we auto-generate a new one.    pub branch_name: Option<String>,    /// SHA of the WIP commit that captures files that were staged (but not yet    /// committed) at the time of archiving. This commit can be empty if the    /// user had no staged files at the time. It sits directly on top of whatever    /// the user's last actual commit was.    pub staged_commit_hash: String,    /// SHA of the WIP commit that captures files that were unstaged (including    /// untracked) at the time of archiving. This commit can be empty if the user    /// had no unstaged files at the time. It sits on top of `staged_commit_hash`.    /// After doing `git reset` past both of these commits, we're back in the state    /// we had before archiving, including what was staged, what was unstaged, and    /// what was committed.    pub unstaged_commit_hash: String,    /// SHA of the commit that HEAD pointed at before we created the two WIP    /// commits during archival. After resetting past the WIP commits during    /// restore, HEAD should land back on this commit. It also serves as a    /// pre-restore sanity check (abort if this commit no longer exists in the    /// repo) and as a fallback target if the WIP resets fail.    pub original_commit_hash: String,} /// The store holds all metadata needed to show threads in the sidebar/the archive.////// Listens to ConversationView events and updates metadata when the root thread changes.pub struct ThreadMetadataStore {

Lines 1–500 of 4273.