3485 lines · 129.1 KB · Rust
RawDownload
#[cfg(test)]mod context_tests; use crate::slash_command_working_set::SlashCommandWorkingSet;use crate::{    prompts::PromptBuilder,    slash_command::{file_command::FileCommandMetadata, SlashCommandLine},    AssistantEdit, AssistantPatch, AssistantPatchStatus, MessageId, MessageStatus,};use anyhow::{anyhow, Context as _, Result};use assistant_slash_command::{    SlashCommandContent, SlashCommandEvent, SlashCommandOutputSection, SlashCommandResult,};use assistant_tool::ToolWorkingSet;use client::{self, proto, telemetry::Telemetry};use clock::ReplicaId;use collections::{HashMap, HashSet};use feature_flags::{FeatureFlagAppExt, ToolUseFeatureFlag};use fs::{Fs, RemoveOptions};use futures::{future::Shared, FutureExt, StreamExt};use gpui::{    AppContext, Context as _, EventEmitter, Model, ModelContext, RenderImage, SharedString,    Subscription, Task,}; use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};use language_model::{    LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,    LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelRequestTool, LanguageModelToolResult, LanguageModelToolUse,    LanguageModelToolUseId, MessageContent, Role, StopReason,};use language_models::{    provider::cloud::{MaxMonthlySpendReachedError, PaymentRequiredError},    report_assistant_event,};use open_ai::Model as OpenAiModel;use paths::contexts_dir;use project::Project;use serde::{Deserialize, Serialize};use smallvec::SmallVec;use std::{    cmp::{max, Ordering},    fmt::Debug,    iter, mem,    ops::Range,    path::{Path, PathBuf},    str::FromStr as _,    sync::Arc,    time::{Duration, Instant},};use telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase};use text::{BufferSnapshot, ToPoint};use util::{post_inc, ResultExt, TryFutureExt};use uuid::Uuid;use workspace::ui::IconName; #[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, Copy, Debug, PartialEq, Eq)]pub enum RequestType {    /// Request a normal chat response from the model.    Chat,    /// Add a preamble to the message, which tells the model to return a structured response that suggests edits.    SuggestEdits,} #[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,    },    SlashCommandStarted {        id: InvokedSlashCommandId,        output_range: Range<language::Anchor>,        name: String,        version: clock::Global,    },    SlashCommandFinished {        id: InvokedSlashCommandId,        timestamp: clock::Lamport,        error_message: Option<String>,        version: clock::Global,    },    SlashCommandOutputSectionAdded {        timestamp: clock::Lamport,        section: 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,                        cache: None,                    },                    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")?,                    ),                    cache: None,                },                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::SlashCommandStarted(message) => {                Ok(Self::SlashCommandStarted {                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(                        message.id.context("invalid id")?,                    )),                    output_range: language::proto::deserialize_anchor_range(                        message.output_range.context("invalid range")?,                    )?,                    name: message.name,                    version: language::proto::deserialize_version(&message.version),                })            }            proto::context_operation::Variant::SlashCommandOutputSectionAdded(message) => {                let section = message.section.context("missing section")?;                Ok(Self::SlashCommandOutputSectionAdded {                    timestamp: language::proto::deserialize_timestamp(                        message.timestamp.context("missing timestamp")?,                    ),                    section: SlashCommandOutputSection {                        range: language::proto::deserialize_anchor_range(                            section.range.context("invalid range")?,                        )?,                        icon: section.icon_name.parse()?,                        label: section.label.into(),                        metadata: section                            .metadata                            .and_then(|metadata| serde_json::from_str(&metadata).log_err()),                    },                    version: language::proto::deserialize_version(&message.version),                })            }            proto::context_operation::Variant::SlashCommandCompleted(message) => {                Ok(Self::SlashCommandFinished {                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(                        message.id.context("invalid id")?,                    )),                    timestamp: language::proto::deserialize_timestamp(                        message.timestamp.context("missing timestamp")?,                    ),                    error_message: message.error_message,                    version: language::proto::deserialize_version(&message.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::SlashCommandStarted {                id,                output_range,                name,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::SlashCommandStarted(                    proto::context_operation::SlashCommandStarted {                        id: Some(language::proto::serialize_timestamp(id.0)),                        output_range: Some(language::proto::serialize_anchor_range(                            output_range.clone(),                        )),                        name: name.clone(),                        version: language::proto::serialize_version(version),                    },                )),            },            Self::SlashCommandOutputSectionAdded {                timestamp,                section,                version,            } => proto::ContextOperation {                variant: Some(                    proto::context_operation::Variant::SlashCommandOutputSectionAdded(                        proto::context_operation::SlashCommandOutputSectionAdded {                            timestamp: Some(language::proto::serialize_timestamp(*timestamp)),                            section: Some({                                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(),                                    metadata: section.metadata.as_ref().and_then(|metadata| {                                        serde_json::to_string(metadata).log_err()                                    }),                                }                            }),                            version: language::proto::serialize_version(version),                        },                    ),                ),            },            Self::SlashCommandFinished {                id,                timestamp,                error_message,                version,            } => proto::ContextOperation {                variant: Some(proto::context_operation::Variant::SlashCommandCompleted(                    proto::context_operation::SlashCommandCompleted {                        id: Some(language::proto::serialize_timestamp(id.0)),                        timestamp: Some(language::proto::serialize_timestamp(*timestamp)),                        error_message: error_message.clone(),                        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::SlashCommandStarted { id, .. } => id.0,            Self::SlashCommandOutputSectionAdded { timestamp, .. }            | Self::SlashCommandFinished { timestamp, .. } => *timestamp,            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::SlashCommandStarted { version, .. }            | Self::SlashCommandOutputSectionAdded { version, .. }            | Self::SlashCommandFinished { version, .. } => version,            Self::BufferOperation(_) => {                panic!("reading the version of a buffer operation is not supported")            }        }    }} #[derive(Debug, Clone)]pub enum ContextEvent {    ShowAssistError(SharedString),    ShowPaymentRequiredError,    ShowMaxMonthlySpendReachedError,    MessagesEdited,    SummaryChanged,    StreamedCompletion,    PatchesUpdated {        removed: Vec<Range<language::Anchor>>,        updated: Vec<Range<language::Anchor>>,    },    InvokedSlashCommandChanged {        command_id: InvokedSlashCommandId,    },    ParsedSlashCommandsUpdated {        removed: Vec<Range<language::Anchor>>,        updated: Vec<ParsedSlashCommand>,    },    SlashCommandOutputSectionAdded {        section: SlashCommandOutputSection<language::Anchor>,    },    UsePendingTools,    ToolFinished {        tool_use_id: LanguageModelToolUseId,        output_range: Range<language::Anchor>,    },    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)]pub enum CacheStatus {    Pending,    Cached,} #[derive(Clone, Debug, Eq, PartialEq)]pub struct MessageCacheMetadata {    pub is_anchor: bool,    pub is_final_anchor: bool,    pub status: CacheStatus,    pub cached_at: clock::Global,} #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]pub struct MessageMetadata {    pub role: Role,    pub status: MessageStatus,    pub(crate) timestamp: clock::Lamport,    #[serde(skip)]    pub cache: Option<MessageCacheMetadata>,} impl From<&Message> for MessageMetadata {    fn from(message: &Message) -> Self {        Self {            role: message.role,            status: message.status.clone(),            timestamp: message.id.0,            cache: message.cache.clone(),        }    }} impl MessageMetadata {    pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {        let result = match &self.cache {            Some(MessageCacheMetadata { cached_at, .. }) => !buffer.has_edits_since_in_range(                &cached_at,                Range {                    start: buffer.anchor_at(range.start, Bias::Right),                    end: buffer.anchor_at(range.end, Bias::Left),                },            ),            _ => false,        };        result    }} #[derive(Clone, Debug)]pub struct Message {    pub offset_range: Range<usize>,    pub index_range: Range<usize>,    pub anchor_range: Range<language::Anchor>,    pub id: MessageId,    pub role: Role,    pub status: MessageStatus,    pub cache: Option<MessageCacheMetadata>,} #[derive(Debug, Clone)]pub enum Content {    Image {        anchor: language::Anchor,        image_id: u64,        render_image: Arc<RenderImage>,        image: Shared<Task<Option<LanguageModelImage>>>,    },    ToolUse {        range: Range<language::Anchor>,        tool_use: LanguageModelToolUse,    },    ToolResult {        range: Range<language::Anchor>,        tool_use_id: LanguageModelToolUseId,    },} impl Content {    fn range(&self) -> Range<language::Anchor> {        match self {            Self::Image { anchor, .. } => *anchor..*anchor,            Self::ToolUse { range, .. } | Self::ToolResult { range, .. } => range.clone(),        }    }     fn cmp(&self, other: &Self, buffer: &BufferSnapshot) -> Ordering {        let self_range = self.range();        let other_range = other.range();        if self_range.end.cmp(&other_range.start, buffer).is_lt() {            Ordering::Less        } else if self_range.start.cmp(&other_range.end, buffer).is_gt() {            Ordering::Greater

Lines 1–500 of 3485.