use crate::{    AssistantContext, ContextEvent, ContextId, ContextOperation, ContextVersion, SavedContext,    SavedContextMetadata,};use anyhow::{Context as _, Result};use assistant_slash_command::{SlashCommandId, SlashCommandWorkingSet};use client::{Client, TypedEnvelope, proto, telemetry::Telemetry};use clock::ReplicaId;use collections::HashMap;use context_server::ContextServerId;use fs::{Fs, RemoveOptions};use futures::StreamExt;use fuzzy::StringMatchCandidate;use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Task, WeakEntity};use language::LanguageRegistry;use paths::contexts_dir;use project::{    Project,    context_server_store::{ContextServerStatus, ContextServerStore},};use prompt_store::PromptBuilder;use regex::Regex;use rpc::AnyProtoClient;use std::sync::LazyLock;use std::{cmp::Reverse, ffi::OsStr, mem, path::Path, sync::Arc, time::Duration};use util::{ResultExt, TryFutureExt}; pub(crate) fn init(client: &AnyProtoClient) {    client.add_entity_message_handler(ContextStore::handle_advertise_contexts);    client.add_entity_request_handler(ContextStore::handle_open_context);    client.add_entity_request_handler(ContextStore::handle_create_context);    client.add_entity_message_handler(ContextStore::handle_update_context);    client.add_entity_request_handler(ContextStore::handle_synchronize_contexts);} #[derive(Clone)]pub struct RemoteContextMetadata {    pub id: ContextId,    pub summary: Option<String>,} pub struct ContextStore {    contexts: Vec<ContextHandle>,    contexts_metadata: Vec<SavedContextMetadata>,    context_server_slash_command_ids: HashMap<ContextServerId, Vec<SlashCommandId>>,    host_contexts: Vec<RemoteContextMetadata>,    fs: Arc<dyn Fs>,    languages: Arc<LanguageRegistry>,    slash_commands: Arc<SlashCommandWorkingSet>,    telemetry: Arc<Telemetry>,    _watch_updates: Task<Option<()>>,    client: Arc<Client>,    project: Entity<Project>,    project_is_shared: bool,    client_subscription: Option<client::Subscription>,    _project_subscriptions: Vec<gpui::Subscription>,    prompt_builder: Arc<PromptBuilder>,} pub enum ContextStoreEvent {    ContextCreated(ContextId),} impl EventEmitter<ContextStoreEvent> for ContextStore {} enum ContextHandle {    Weak(WeakEntity<AssistantContext>),    Strong(Entity<AssistantContext>),} impl ContextHandle {    fn upgrade(&self) -> Option<Entity<AssistantContext>> {        match self {            ContextHandle::Weak(weak) => weak.upgrade(),            ContextHandle::Strong(strong) => Some(strong.clone()),        }    }     fn downgrade(&self) -> WeakEntity<AssistantContext> {        match self {            ContextHandle::Weak(weak) => weak.clone(),            ContextHandle::Strong(strong) => strong.downgrade(),        }    }} impl ContextStore {    pub fn new(        project: Entity<Project>,        prompt_builder: Arc<PromptBuilder>,        slash_commands: Arc<SlashCommandWorkingSet>,        cx: &mut App,    ) -> Task<Result<Entity<Self>>> {        let fs = project.read(cx).fs().clone();        let languages = project.read(cx).languages().clone();        let telemetry = project.read(cx).client().telemetry().clone();        cx.spawn(async move |cx| {            const CONTEXT_WATCH_DURATION: Duration = Duration::from_millis(100);            let (mut events, _) = fs.watch(contexts_dir(), CONTEXT_WATCH_DURATION).await;             let this = cx.new(|cx: &mut Context<Self>| {                let mut this = Self {                    contexts: Vec::new(),                    contexts_metadata: Vec::new(),                    context_server_slash_command_ids: HashMap::default(),                    host_contexts: Vec::new(),                    fs,                    languages,                    slash_commands,                    telemetry,                    _watch_updates: cx.spawn(async move |this, cx| {                        async move {                            while events.next().await.is_some() {                                this.update(cx, |this, cx| this.reload(cx))?.await.log_err();                            }                            anyhow::Ok(())                        }                        .log_err()                        .await                    }),                    client_subscription: None,                    _project_subscriptions: vec![                        cx.subscribe(&project, Self::handle_project_event),                    ],                    project_is_shared: false,                    client: project.read(cx).client(),                    project: project.clone(),                    prompt_builder,                };                this.handle_project_shared(project.clone(), cx);                this.synchronize_contexts(cx);                this.register_context_server_handlers(cx);                this.reload(cx).detach_and_log_err(cx);                this            })?;             Ok(this)        })    }     async fn handle_advertise_contexts(        this: Entity<Self>,        envelope: TypedEnvelope<proto::AdvertiseContexts>,        mut cx: AsyncApp,    ) -> Result<()> {        this.update(&mut cx, |this, cx| {            this.host_contexts = envelope                .payload                .contexts                .into_iter()                .map(|context| RemoteContextMetadata {                    id: ContextId::from_proto(context.context_id),                    summary: context.summary,                })                .collect();            cx.notify();        })    }     async fn handle_open_context(        this: Entity<Self>,        envelope: TypedEnvelope<proto::OpenContext>,        mut cx: AsyncApp,    ) -> Result<proto::OpenContextResponse> {        let context_id = ContextId::from_proto(envelope.payload.context_id);        let operations = this.update(&mut cx, |this, cx| {            anyhow::ensure!(                !this.project.read(cx).is_via_collab(),                "only the host contexts can be opened"            );             let context = this                .loaded_context_for_id(&context_id, cx)                .context("context not found")?;            anyhow::ensure!(                context.read(cx).replica_id() == ReplicaId::default(),                "context must be opened via the host"            );             anyhow::Ok(                context                    .read(cx)                    .serialize_ops(&ContextVersion::default(), cx),            )        })??;        let operations = operations.await;        Ok(proto::OpenContextResponse {            context: Some(proto::Context { operations }),        })    }     async fn handle_create_context(        this: Entity<Self>,        _: TypedEnvelope<proto::CreateContext>,        mut cx: AsyncApp,    ) -> Result<proto::CreateContextResponse> {        let (context_id, operations) = this.update(&mut cx, |this, cx| {            anyhow::ensure!(                !this.project.read(cx).is_via_collab(),                "can only create contexts as the host"            );             let context = this.create(cx);            let context_id = context.read(cx).id().clone();            cx.emit(ContextStoreEvent::ContextCreated(context_id.clone()));             anyhow::Ok((                context_id,                context                    .read(cx)                    .serialize_ops(&ContextVersion::default(), cx),            ))        })??;        let operations = operations.await;        Ok(proto::CreateContextResponse {            context_id: context_id.to_proto(),            context: Some(proto::Context { operations }),        })    }     async fn handle_update_context(        this: Entity<Self>,        envelope: TypedEnvelope<proto::UpdateContext>,        mut cx: AsyncApp,    ) -> Result<()> {        this.update(&mut cx, |this, cx| {            let context_id = ContextId::from_proto(envelope.payload.context_id);            if let Some(context) = this.loaded_context_for_id(&context_id, cx) {                let operation_proto = envelope.payload.operation.context("invalid operation")?;                let operation = ContextOperation::from_proto(operation_proto)?;                context.update(cx, |context, cx| context.apply_ops([operation], cx));            }            Ok(())        })?    }     async fn handle_synchronize_contexts(        this: Entity<Self>,        envelope: TypedEnvelope<proto::SynchronizeContexts>,        mut cx: AsyncApp,    ) -> Result<proto::SynchronizeContextsResponse> {        this.update(&mut cx, |this, cx| {            anyhow::ensure!(                !this.project.read(cx).is_via_collab(),                "only the host can synchronize contexts"            );             let mut local_versions = Vec::new();            for remote_version_proto in envelope.payload.contexts {                let remote_version = ContextVersion::from_proto(&remote_version_proto);                let context_id = ContextId::from_proto(remote_version_proto.context_id);                if let Some(context) = this.loaded_context_for_id(&context_id, cx) {                    let context = context.read(cx);                    let operations = context.serialize_ops(&remote_version, cx);                    local_versions.push(context.version(cx).to_proto(context_id.clone()));                    let client = this.client.clone();                    let project_id = envelope.payload.project_id;                    cx.background_spawn(async move {                        let operations = operations.await;                        for operation in operations {                            client.send(proto::UpdateContext {                                project_id,                                context_id: context_id.to_proto(),                                operation: Some(operation),                            })?;                        }                        anyhow::Ok(())                    })                    .detach_and_log_err(cx);                }            }             this.advertise_contexts(cx);             anyhow::Ok(proto::SynchronizeContextsResponse {                contexts: local_versions,            })        })?    }     fn handle_project_shared(&mut self, _: Entity<Project>, cx: &mut Context<Self>) {        let is_shared = self.project.read(cx).is_shared();        let was_shared = mem::replace(&mut self.project_is_shared, is_shared);        if is_shared == was_shared {            return;        }         if is_shared {            self.contexts.retain_mut(|context| {                if let Some(strong_context) = context.upgrade() {                    *context = ContextHandle::Strong(strong_context);                    true                } else {                    false                }            });            let remote_id = self.project.read(cx).remote_id().unwrap();            self.client_subscription = self                .client                .subscribe_to_entity(remote_id)                .log_err()                .map(|subscription| subscription.set_entity(&cx.entity(), &mut cx.to_async()));            self.advertise_contexts(cx);        } else {            self.client_subscription = None;        }    }     fn handle_project_event(        &mut self,        project: Entity<Project>,        event: &project::Event,        cx: &mut Context<Self>,    ) {        match event {            project::Event::RemoteIdChanged(_) => {                self.handle_project_shared(project, cx);            }            project::Event::Reshared => {                self.advertise_contexts(cx);            }            project::Event::HostReshared | project::Event::Rejoined => {                self.synchronize_contexts(cx);            }            project::Event::DisconnectedFromHost => {                self.contexts.retain_mut(|context| {                    if let Some(strong_context) = context.upgrade() {                        *context = ContextHandle::Weak(context.downgrade());                        strong_context.update(cx, |context, cx| {                            if context.replica_id() != ReplicaId::default() {                                context.set_capability(language::Capability::ReadOnly, cx);                            }                        });                        true                    } else {                        false                    }                });                self.host_contexts.clear();                cx.notify();            }            _ => {}        }    }     pub fn unordered_contexts(&self) -> impl Iterator<Item = &SavedContextMetadata> {        self.contexts_metadata.iter()    }     pub fn reverse_chronological_contexts(&self) -> Vec<SavedContextMetadata> {        let mut contexts = self.contexts_metadata.iter().cloned().collect::<Vec<_>>();        contexts.sort_unstable_by_key(|thread| std::cmp::Reverse(thread.mtime));        contexts    }     pub fn create(&mut self, cx: &mut Context<Self>) -> Entity<AssistantContext> {        let context = cx.new(|cx| {            AssistantContext::local(                self.languages.clone(),                Some(self.project.clone()),                Some(self.telemetry.clone()),                self.prompt_builder.clone(),                self.slash_commands.clone(),                cx,            )        });        self.register_context(&context, cx);        context    }     pub fn create_remote_context(        &mut self,        cx: &mut Context<Self>,    ) -> Task<Result<Entity<AssistantContext>>> {        let project = self.project.read(cx);        let Some(project_id) = project.remote_id() else {            return Task::ready(Err(anyhow::anyhow!("project was not remote")));        };         let replica_id = project.replica_id();        let capability = project.capability();        let language_registry = self.languages.clone();        let project = self.project.clone();        let telemetry = self.telemetry.clone();        let prompt_builder = self.prompt_builder.clone();        let slash_commands = self.slash_commands.clone();        let request = self.client.request(proto::CreateContext { project_id });        cx.spawn(async move |this, cx| {            let response = request.await?;            let context_id = ContextId::from_proto(response.context_id);            let context_proto = response.context.context("invalid context")?;            let context = cx.new(|cx| {                AssistantContext::new(                    context_id.clone(),                    replica_id,                    capability,                    language_registry,                    prompt_builder,                    slash_commands,                    Some(project),                    Some(telemetry),                    cx,                )            })?;            let operations = cx                .background_spawn(async move {                    context_proto                        .operations                        .into_iter()                        .map(ContextOperation::from_proto)                        .collect::<Result<Vec<_>>>()                })                .await?;            context.update(cx, |context, cx| context.apply_ops(operations, cx))?;            this.update(cx, |this, cx| {                if let Some(existing_context) = this.loaded_context_for_id(&context_id, cx) {                    existing_context                } else {                    this.register_context(&context, cx);                    this.synchronize_contexts(cx);                    context                }            })        })    }     pub fn open_local_context(        &mut self,        path: Arc<Path>,        cx: &Context<Self>,    ) -> Task<Result<Entity<AssistantContext>>> {        if let Some(existing_context) = self.loaded_context_for_path(&path, cx) {            return Task::ready(Ok(existing_context));        }         let fs = self.fs.clone();        let languages = self.languages.clone();        let project = self.project.clone();        let telemetry = self.telemetry.clone();        let load = cx.background_spawn({            let path = path.clone();            async move {                let saved_context = fs.load(&path).await?;                SavedContext::from_json(&saved_context)            }        });        let prompt_builder = self.prompt_builder.clone();        let slash_commands = self.slash_commands.clone();         cx.spawn(async move |this, cx| {            let saved_context = load.await?;            let context = cx.new(|cx| {                AssistantContext::deserialize(                    saved_context,                    path.clone(),                    languages,                    prompt_builder,                    slash_commands,                    Some(project),                    Some(telemetry),                    cx,                )            })?;            this.update(cx, |this, cx| {                if let Some(existing_context) = this.loaded_context_for_path(&path, cx) {                    existing_context                } else {                    this.register_context(&context, cx);                    context                }            })        })    }     pub fn delete_local_context(        &mut self,        path: Arc<Path>,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let fs = self.fs.clone();         cx.spawn(async move |this, cx| {            fs.remove_file(                &path,                RemoveOptions {                    recursive: false,                    ignore_if_not_exists: true,                },            )            .await?;             this.update(cx, |this, cx| {                this.contexts.retain(|context| {                    context                        .upgrade()                        .and_then(|context| context.read(cx).path())                        != Some(&path)                });                this.contexts_metadata                    .retain(|context| context.path.as_ref() != path.as_ref());

Lines 1–500 of 887. Load more