3009 lines · 109.5 KB · Rust
RawDownload
use crate::{    slash_command::SlashCommandLine, CompletionProvider, LanguageModelRequest,    LanguageModelRequestMessage, MessageId, MessageStatus, Role,};use anyhow::{anyhow, Context as _, Result};use assistant_slash_command::{    SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry,};use client::{proto, telemetry::Telemetry};use clock::ReplicaId;use collections::{HashMap, HashSet};use fs::Fs;use futures::{future::Shared, FutureExt, StreamExt};use gpui::{AppContext, Context as _, EventEmitter, Model, ModelContext, Subscription, Task};use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};use open_ai::Model as OpenAiModel;use paths::contexts_dir;use serde::{Deserialize, Serialize};use std::{    cmp::Ordering,    iter, mem,    ops::Range,    path::{Path, PathBuf},    sync::Arc,    time::{Duration, Instant},};use telemetry_events::AssistantKind;use ui::SharedString;use util::{post_inc, TryFutureExt};use uuid::Uuid; #[derive(Clone, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]pub struct ContextId(String); impl ContextId {    pub fn new() -> Self {        Self(Uuid::new_v4().to_string())    }     pub fn from_proto(id: String) -> Self {        Self(id)    }     pub fn to_proto(&self) -> String {        self.0.clone()    }} #[derive(Clone, Debug)]pub enum ContextOperation {    InsertMessage {        anchor: MessageAnchor,        metadata: MessageMetadata,        version: clock::Global,    },    UpdateMessage {        message_id: MessageId,        metadata: MessageMetadata,        version: clock::Global,    },    UpdateSummary {        summary: ContextSummary,        version: clock::Global,    },    SlashCommandFinished {        id: SlashCommandId,        output_range: Range<language::Anchor>,        sections: Vec<SlashCommandOutputSection<language::Anchor>>,        version: clock::Global,    },    BufferOperation(language::Operation),} impl ContextOperation {    pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {        match op.variant.context("invalid variant")? {            proto::context_operation::Variant::InsertMessage(insert) => {                let message = insert.message.context("invalid message")?;                let id = MessageId(language::proto::deserialize_timestamp(                    message.id.context("invalid id")?,                ));                Ok(Self::InsertMessage {                    anchor: MessageAnchor {                        id,                        start: language::proto::deserialize_anchor(                            message.start.context("invalid anchor")?,                        )                        .context("invalid anchor")?,                    },                    metadata: MessageMetadata {                        role: Role::from_proto(message.role),                        status: MessageStatus::from_proto(                            message.status.context("invalid status")?,                        ),                        timestamp: id.0,                    },                    version: language::proto::deserialize_version(&insert.version),                })            }            proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {                message_id: MessageId(language::proto::deserialize_timestamp(                    update.message_id.context("invalid message id")?,                )),                metadata: MessageMetadata {                    role: Role::from_proto(update.role),                    status: MessageStatus::from_proto(update.status.context("invalid status")?),                    timestamp: language::proto::deserialize_timestamp(                        update.timestamp.context("invalid timestamp")?,                    ),                },                version: language::proto::deserialize_version(&update.version),            }),            proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {                summary: ContextSummary {                    text: update.summary,                    done: update.done,                    timestamp: language::proto::deserialize_timestamp(                        update.timestamp.context("invalid timestamp")?,                    ),                },                version: language::proto::deserialize_version(&update.version),            }),            proto::context_operation::Variant::SlashCommandFinished(finished) => {                Ok(Self::SlashCommandFinished {                    id: SlashCommandId(language::proto::deserialize_timestamp(                        finished.id.context("invalid id")?,                    )),                    output_range: language::proto::deserialize_anchor_range(                        finished.output_range.context("invalid range")?,                    )?,                    sections: finished                        .sections                        .into_iter()                        .map(|section| {                            Ok(SlashCommandOutputSection {                                range: language::proto::deserialize_anchor_range(                                    section.range.context("invalid range")?,                                )?,                                icon: section.icon_name.parse()?,                                label: section.label.into(),                            })                        })                        .collect::<Result<Vec<_>>>()?,                    version: language::proto::deserialize_version(&finished.version),                })            }            proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(                language::proto::deserialize_operation(                    op.operation.context("invalid buffer operation")?,                )?,            )),        }    }     pub fn to_proto(&self) -> proto::ContextOperation {        match self {            Self::InsertMessage {                anchor,                metadata,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::InsertMessage(                    proto::context_operation::InsertMessage {                        message: Some(proto::ContextMessage {                            id: Some(language::proto::serialize_timestamp(anchor.id.0)),                            start: Some(language::proto::serialize_anchor(&anchor.start)),                            role: metadata.role.to_proto() as i32,                            status: Some(metadata.status.to_proto()),                        }),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::UpdateMessage {                message_id,                metadata,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::UpdateMessage(                    proto::context_operation::UpdateMessage {                        message_id: Some(language::proto::serialize_timestamp(message_id.0)),                        role: metadata.role.to_proto() as i32,                        status: Some(metadata.status.to_proto()),                        timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::UpdateSummary { summary, version } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::UpdateSummary(                    proto::context_operation::UpdateSummary {                        summary: summary.text.clone(),                        done: summary.done,                        timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::SlashCommandFinished {                id,                output_range,                sections,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::SlashCommandFinished(                    proto::context_operation::SlashCommandFinished {                        id: Some(language::proto::serialize_timestamp(id.0)),                        output_range: Some(language::proto::serialize_anchor_range(                            output_range.clone(),                        )),                        sections: sections                            .iter()                            .map(|section| {                                let icon_name: &'static str = section.icon.into();                                proto::SlashCommandOutputSection {                                    range: Some(language::proto::serialize_anchor_range(                                        section.range.clone(),                                    )),                                    icon_name: icon_name.to_string(),                                    label: section.label.to_string(),                                }                            })                            .collect(),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::BufferOperation(operation) => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::BufferOperation(                    proto::context_operation::BufferOperation {                        operation: Some(language::proto::serialize_operation(operation)),                    },                )),            },        }    }     fn timestamp(&self) -> clock::Lamport {        match self {            Self::InsertMessage { anchor, .. } => anchor.id.0,            Self::UpdateMessage { metadata, .. } => metadata.timestamp,            Self::UpdateSummary { summary, .. } => summary.timestamp,            Self::SlashCommandFinished { id, .. } => id.0,            Self::BufferOperation(_) => {                panic!("reading the timestamp of a buffer operation is not supported")            }        }    }     /// Returns the current version of the context operation.    pub fn version(&self) -> &clock::Global {        match self {            Self::InsertMessage { version, .. }            | Self::UpdateMessage { version, .. }            | Self::UpdateSummary { version, .. }            | Self::SlashCommandFinished { version, .. } => version,            Self::BufferOperation(_) => {                panic!("reading the version of a buffer operation is not supported")            }        }    }} #[derive(Clone)]pub enum ContextEvent {    MessagesEdited,    SummaryChanged,    EditSuggestionsChanged,    StreamedCompletion,    PendingSlashCommandsUpdated {        removed: Vec<Range<language::Anchor>>,        updated: Vec<PendingSlashCommand>,    },    SlashCommandFinished {        output_range: Range<language::Anchor>,        sections: Vec<SlashCommandOutputSection<language::Anchor>>,        run_commands_in_output: bool,    },    Operation(ContextOperation),} #[derive(Clone, Default, Debug)]pub struct ContextSummary {    pub text: String,    done: bool,    timestamp: clock::Lamport,} #[derive(Clone, Debug, Eq, PartialEq)]pub struct MessageAnchor {    pub id: MessageId,    pub start: language::Anchor,} #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]pub struct MessageMetadata {    pub role: Role,    status: MessageStatus,    timestamp: clock::Lamport,} #[derive(Clone, Debug, PartialEq, Eq)]pub struct Message {    pub offset_range: Range<usize>,    pub index_range: Range<usize>,    pub id: MessageId,    pub anchor: language::Anchor,    pub role: Role,    pub status: MessageStatus,} impl Message {    fn to_request_message(&self, buffer: &Buffer) -> LanguageModelRequestMessage {        LanguageModelRequestMessage {            role: self.role,            content: buffer.text_for_range(self.offset_range.clone()).collect(),        }    }} struct PendingCompletion {    id: usize,    _task: Task<()>,} #[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]pub struct SlashCommandId(clock::Lamport); pub struct Context {    id: ContextId,    timestamp: clock::Lamport,    version: clock::Global,    pending_ops: Vec<ContextOperation>,    operations: Vec<ContextOperation>,    buffer: Model<Buffer>,    edit_suggestions: Vec<EditSuggestion>,    pending_slash_commands: Vec<PendingSlashCommand>,    edits_since_last_slash_command_parse: language::Subscription,    finished_slash_commands: HashSet<SlashCommandId>,    slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,    message_anchors: Vec<MessageAnchor>,    messages_metadata: HashMap<MessageId, MessageMetadata>,    summary: Option<ContextSummary>,    pending_summary: Task<Option<()>>,    completion_count: usize,    pending_completions: Vec<PendingCompletion>,    token_count: Option<usize>,    pending_token_count: Task<Option<()>>,    pending_edit_suggestion_parse: Option<Task<()>>,    pending_save: Task<Result<()>>,    path: Option<PathBuf>,    _subscriptions: Vec<Subscription>,    telemetry: Option<Arc<Telemetry>>,    language_registry: Arc<LanguageRegistry>,} impl EventEmitter<ContextEvent> for Context {} impl Context {    pub fn local(        language_registry: Arc<LanguageRegistry>,        telemetry: Option<Arc<Telemetry>>,        cx: &mut ModelContext<Self>,    ) -> Self {        Self::new(            ContextId::new(),            ReplicaId::default(),            language::Capability::ReadWrite,            language_registry,            telemetry,            cx,        )    }     pub fn new(        id: ContextId,        replica_id: ReplicaId,        capability: language::Capability,        language_registry: Arc<LanguageRegistry>,        telemetry: Option<Arc<Telemetry>>,        cx: &mut ModelContext<Self>,    ) -> Self {        let buffer = cx.new_model(|_cx| {            let mut buffer = Buffer::remote(                language::BufferId::new(1).unwrap(),                replica_id,                capability,                "",            );            buffer.set_language_registry(language_registry.clone());            buffer        });        let edits_since_last_slash_command_parse =            buffer.update(cx, |buffer, _| buffer.subscribe());        let mut this = Self {            id,            timestamp: clock::Lamport::new(replica_id),            version: clock::Global::new(),            pending_ops: Vec::new(),            operations: Vec::new(),            message_anchors: Default::default(),            messages_metadata: Default::default(),            edit_suggestions: Vec::new(),            pending_slash_commands: Vec::new(),            finished_slash_commands: HashSet::default(),            slash_command_output_sections: Vec::new(),            edits_since_last_slash_command_parse,            summary: None,            pending_summary: Task::ready(None),            completion_count: Default::default(),            pending_completions: Default::default(),            token_count: None,            pending_token_count: Task::ready(None),            pending_edit_suggestion_parse: None,            _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],            pending_save: Task::ready(Ok(())),            path: None,            buffer,            telemetry,            language_registry,        };         let first_message_id = MessageId(clock::Lamport {            replica_id: 0,            value: 0,        });        let message = MessageAnchor {            id: first_message_id,            start: language::Anchor::MIN,        };        this.messages_metadata.insert(            first_message_id,            MessageMetadata {                role: Role::User,                status: MessageStatus::Done,                timestamp: first_message_id.0,            },        );        this.message_anchors.push(message);         this.set_language(cx);        this.count_remaining_tokens(cx);        this    }     fn serialize(&self, cx: &AppContext) -> SavedContext {        let buffer = self.buffer.read(cx);        SavedContext {            id: Some(self.id.clone()),            zed: "context".into(),            version: SavedContext::VERSION.into(),            text: buffer.text(),            messages: self                .messages(cx)                .map(|message| SavedMessage {                    id: message.id,                    start: message.offset_range.start,                    metadata: self.messages_metadata[&message.id].clone(),                })                .collect(),            summary: self                .summary                .as_ref()                .map(|summary| summary.text.clone())                .unwrap_or_default(),            slash_command_output_sections: self                .slash_command_output_sections                .iter()                .filter_map(|section| {                    let range = section.range.to_offset(buffer);                    if section.range.start.is_valid(buffer) && !range.is_empty() {                        Some(assistant_slash_command::SlashCommandOutputSection {                            range,                            icon: section.icon,                            label: section.label.clone(),                        })                    } else {                        None                    }                })                .collect(),        }    }     #[allow(clippy::too_many_arguments)]    pub fn deserialize(        saved_context: SavedContext,        path: PathBuf,        language_registry: Arc<LanguageRegistry>,        telemetry: Option<Arc<Telemetry>>,        cx: &mut ModelContext<Self>,    ) -> Self {        let id = saved_context.id.clone().unwrap_or_else(|| ContextId::new());        let mut this = Self::new(            id,            ReplicaId::default(),            language::Capability::ReadWrite,            language_registry,            telemetry,            cx,

Lines 1–500 of 3009.