context.rs

   1#[cfg(test)]
   2mod context_tests;
   3
   4use crate::{
   5    prompts::PromptBuilder, slash_command::SlashCommandLine, AssistantEdit, AssistantPatch,
   6    AssistantPatchStatus, MessageId, MessageStatus,
   7};
   8use anyhow::{anyhow, Context as _, Result};
   9use assistant_slash_command::{
  10    SlashCommandOutput, SlashCommandOutputSection, SlashCommandRegistry, SlashCommandResult,
  11};
  12use assistant_tool::ToolRegistry;
  13use client::{self, proto, telemetry::Telemetry};
  14use clock::ReplicaId;
  15use collections::{HashMap, HashSet};
  16use feature_flags::{FeatureFlag, FeatureFlagAppExt};
  17use fs::{Fs, RemoveOptions};
  18use futures::{future::Shared, FutureExt, StreamExt};
  19use gpui::{
  20    AppContext, Context as _, EventEmitter, Model, ModelContext, RenderImage, SharedString,
  21    Subscription, Task,
  22};
  23
  24use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};
  25use language_model::{
  26    provider::cloud::{MaxMonthlySpendReachedError, PaymentRequiredError},
  27    LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,
  28    LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
  29    LanguageModelRequestTool, LanguageModelToolResult, LanguageModelToolUse, MessageContent, Role,
  30    StopReason,
  31};
  32use open_ai::Model as OpenAiModel;
  33use paths::contexts_dir;
  34use project::Project;
  35use serde::{Deserialize, Serialize};
  36use smallvec::SmallVec;
  37use std::{
  38    cmp::{max, Ordering},
  39    fmt::Debug,
  40    iter, mem,
  41    ops::Range,
  42    path::{Path, PathBuf},
  43    str::FromStr as _,
  44    sync::Arc,
  45    time::{Duration, Instant},
  46};
  47use telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase};
  48use text::BufferSnapshot;
  49use util::{post_inc, ResultExt, TryFutureExt};
  50use uuid::Uuid;
  51
  52#[derive(Clone, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
  53pub struct ContextId(String);
  54
  55impl ContextId {
  56    pub fn new() -> Self {
  57        Self(Uuid::new_v4().to_string())
  58    }
  59
  60    pub fn from_proto(id: String) -> Self {
  61        Self(id)
  62    }
  63
  64    pub fn to_proto(&self) -> String {
  65        self.0.clone()
  66    }
  67}
  68
  69#[derive(Clone, Copy, Debug)]
  70pub enum RequestType {
  71    /// Request a normal chat response from the model.
  72    Chat,
  73    /// Add a preamble to the message, which tells the model to return a structured response that suggests edits.
  74    SuggestEdits,
  75}
  76
  77#[derive(Clone, Debug)]
  78pub enum ContextOperation {
  79    InsertMessage {
  80        anchor: MessageAnchor,
  81        metadata: MessageMetadata,
  82        version: clock::Global,
  83    },
  84    UpdateMessage {
  85        message_id: MessageId,
  86        metadata: MessageMetadata,
  87        version: clock::Global,
  88    },
  89    UpdateSummary {
  90        summary: ContextSummary,
  91        version: clock::Global,
  92    },
  93    SlashCommandFinished {
  94        id: SlashCommandId,
  95        output_range: Range<language::Anchor>,
  96        sections: Vec<SlashCommandOutputSection<language::Anchor>>,
  97        version: clock::Global,
  98    },
  99    BufferOperation(language::Operation),
 100}
 101
 102impl ContextOperation {
 103    pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {
 104        match op.variant.context("invalid variant")? {
 105            proto::context_operation::Variant::InsertMessage(insert) => {
 106                let message = insert.message.context("invalid message")?;
 107                let id = MessageId(language::proto::deserialize_timestamp(
 108                    message.id.context("invalid id")?,
 109                ));
 110                Ok(Self::InsertMessage {
 111                    anchor: MessageAnchor {
 112                        id,
 113                        start: language::proto::deserialize_anchor(
 114                            message.start.context("invalid anchor")?,
 115                        )
 116                        .context("invalid anchor")?,
 117                    },
 118                    metadata: MessageMetadata {
 119                        role: Role::from_proto(message.role),
 120                        status: MessageStatus::from_proto(
 121                            message.status.context("invalid status")?,
 122                        ),
 123                        timestamp: id.0,
 124                        cache: None,
 125                    },
 126                    version: language::proto::deserialize_version(&insert.version),
 127                })
 128            }
 129            proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {
 130                message_id: MessageId(language::proto::deserialize_timestamp(
 131                    update.message_id.context("invalid message id")?,
 132                )),
 133                metadata: MessageMetadata {
 134                    role: Role::from_proto(update.role),
 135                    status: MessageStatus::from_proto(update.status.context("invalid status")?),
 136                    timestamp: language::proto::deserialize_timestamp(
 137                        update.timestamp.context("invalid timestamp")?,
 138                    ),
 139                    cache: None,
 140                },
 141                version: language::proto::deserialize_version(&update.version),
 142            }),
 143            proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {
 144                summary: ContextSummary {
 145                    text: update.summary,
 146                    done: update.done,
 147                    timestamp: language::proto::deserialize_timestamp(
 148                        update.timestamp.context("invalid timestamp")?,
 149                    ),
 150                },
 151                version: language::proto::deserialize_version(&update.version),
 152            }),
 153            proto::context_operation::Variant::SlashCommandFinished(finished) => {
 154                Ok(Self::SlashCommandFinished {
 155                    id: SlashCommandId(language::proto::deserialize_timestamp(
 156                        finished.id.context("invalid id")?,
 157                    )),
 158                    output_range: language::proto::deserialize_anchor_range(
 159                        finished.output_range.context("invalid range")?,
 160                    )?,
 161                    sections: finished
 162                        .sections
 163                        .into_iter()
 164                        .map(|section| {
 165                            Ok(SlashCommandOutputSection {
 166                                range: language::proto::deserialize_anchor_range(
 167                                    section.range.context("invalid range")?,
 168                                )?,
 169                                icon: section.icon_name.parse()?,
 170                                label: section.label.into(),
 171                                metadata: section
 172                                    .metadata
 173                                    .and_then(|metadata| serde_json::from_str(&metadata).log_err()),
 174                            })
 175                        })
 176                        .collect::<Result<Vec<_>>>()?,
 177                    version: language::proto::deserialize_version(&finished.version),
 178                })
 179            }
 180            proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(
 181                language::proto::deserialize_operation(
 182                    op.operation.context("invalid buffer operation")?,
 183                )?,
 184            )),
 185        }
 186    }
 187
 188    pub fn to_proto(&self) -> proto::ContextOperation {
 189        match self {
 190            Self::InsertMessage {
 191                anchor,
 192                metadata,
 193                version,
 194            } => proto::ContextOperation {
 195                variant: Some(proto::context_operation::Variant::InsertMessage(
 196                    proto::context_operation::InsertMessage {
 197                        message: Some(proto::ContextMessage {
 198                            id: Some(language::proto::serialize_timestamp(anchor.id.0)),
 199                            start: Some(language::proto::serialize_anchor(&anchor.start)),
 200                            role: metadata.role.to_proto() as i32,
 201                            status: Some(metadata.status.to_proto()),
 202                        }),
 203                        version: language::proto::serialize_version(version),
 204                    },
 205                )),
 206            },
 207            Self::UpdateMessage {
 208                message_id,
 209                metadata,
 210                version,
 211            } => proto::ContextOperation {
 212                variant: Some(proto::context_operation::Variant::UpdateMessage(
 213                    proto::context_operation::UpdateMessage {
 214                        message_id: Some(language::proto::serialize_timestamp(message_id.0)),
 215                        role: metadata.role.to_proto() as i32,
 216                        status: Some(metadata.status.to_proto()),
 217                        timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),
 218                        version: language::proto::serialize_version(version),
 219                    },
 220                )),
 221            },
 222            Self::UpdateSummary { summary, version } => proto::ContextOperation {
 223                variant: Some(proto::context_operation::Variant::UpdateSummary(
 224                    proto::context_operation::UpdateSummary {
 225                        summary: summary.text.clone(),
 226                        done: summary.done,
 227                        timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),
 228                        version: language::proto::serialize_version(version),
 229                    },
 230                )),
 231            },
 232            Self::SlashCommandFinished {
 233                id,
 234                output_range,
 235                sections,
 236                version,
 237            } => proto::ContextOperation {
 238                variant: Some(proto::context_operation::Variant::SlashCommandFinished(
 239                    proto::context_operation::SlashCommandFinished {
 240                        id: Some(language::proto::serialize_timestamp(id.0)),
 241                        output_range: Some(language::proto::serialize_anchor_range(
 242                            output_range.clone(),
 243                        )),
 244                        sections: sections
 245                            .iter()
 246                            .map(|section| {
 247                                let icon_name: &'static str = section.icon.into();
 248                                proto::SlashCommandOutputSection {
 249                                    range: Some(language::proto::serialize_anchor_range(
 250                                        section.range.clone(),
 251                                    )),
 252                                    icon_name: icon_name.to_string(),
 253                                    label: section.label.to_string(),
 254                                    metadata: section.metadata.as_ref().and_then(|metadata| {
 255                                        serde_json::to_string(metadata).log_err()
 256                                    }),
 257                                }
 258                            })
 259                            .collect(),
 260                        version: language::proto::serialize_version(version),
 261                    },
 262                )),
 263            },
 264            Self::BufferOperation(operation) => proto::ContextOperation {
 265                variant: Some(proto::context_operation::Variant::BufferOperation(
 266                    proto::context_operation::BufferOperation {
 267                        operation: Some(language::proto::serialize_operation(operation)),
 268                    },
 269                )),
 270            },
 271        }
 272    }
 273
 274    fn timestamp(&self) -> clock::Lamport {
 275        match self {
 276            Self::InsertMessage { anchor, .. } => anchor.id.0,
 277            Self::UpdateMessage { metadata, .. } => metadata.timestamp,
 278            Self::UpdateSummary { summary, .. } => summary.timestamp,
 279            Self::SlashCommandFinished { id, .. } => id.0,
 280            Self::BufferOperation(_) => {
 281                panic!("reading the timestamp of a buffer operation is not supported")
 282            }
 283        }
 284    }
 285
 286    /// Returns the current version of the context operation.
 287    pub fn version(&self) -> &clock::Global {
 288        match self {
 289            Self::InsertMessage { version, .. }
 290            | Self::UpdateMessage { version, .. }
 291            | Self::UpdateSummary { version, .. }
 292            | Self::SlashCommandFinished { version, .. } => version,
 293            Self::BufferOperation(_) => {
 294                panic!("reading the version of a buffer operation is not supported")
 295            }
 296        }
 297    }
 298}
 299
 300#[derive(Debug, Clone)]
 301pub enum ContextEvent {
 302    ShowAssistError(SharedString),
 303    ShowPaymentRequiredError,
 304    ShowMaxMonthlySpendReachedError,
 305    MessagesEdited,
 306    SummaryChanged,
 307    StreamedCompletion,
 308    PatchesUpdated {
 309        removed: Vec<Range<language::Anchor>>,
 310        updated: Vec<Range<language::Anchor>>,
 311    },
 312    PendingSlashCommandsUpdated {
 313        removed: Vec<Range<language::Anchor>>,
 314        updated: Vec<PendingSlashCommand>,
 315    },
 316    SlashCommandFinished {
 317        output_range: Range<language::Anchor>,
 318        sections: Vec<SlashCommandOutputSection<language::Anchor>>,
 319        run_commands_in_output: bool,
 320        expand_result: bool,
 321    },
 322    UsePendingTools,
 323    ToolFinished {
 324        tool_use_id: Arc<str>,
 325        output_range: Range<language::Anchor>,
 326    },
 327    Operation(ContextOperation),
 328}
 329
 330#[derive(Clone, Default, Debug)]
 331pub struct ContextSummary {
 332    pub text: String,
 333    done: bool,
 334    timestamp: clock::Lamport,
 335}
 336
 337#[derive(Clone, Debug, Eq, PartialEq)]
 338pub struct MessageAnchor {
 339    pub id: MessageId,
 340    pub start: language::Anchor,
 341}
 342
 343#[derive(Clone, Debug, Eq, PartialEq)]
 344pub enum CacheStatus {
 345    Pending,
 346    Cached,
 347}
 348
 349#[derive(Clone, Debug, Eq, PartialEq)]
 350pub struct MessageCacheMetadata {
 351    pub is_anchor: bool,
 352    pub is_final_anchor: bool,
 353    pub status: CacheStatus,
 354    pub cached_at: clock::Global,
 355}
 356
 357#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 358pub struct MessageMetadata {
 359    pub role: Role,
 360    pub status: MessageStatus,
 361    pub(crate) timestamp: clock::Lamport,
 362    #[serde(skip)]
 363    pub cache: Option<MessageCacheMetadata>,
 364}
 365
 366impl From<&Message> for MessageMetadata {
 367    fn from(message: &Message) -> Self {
 368        Self {
 369            role: message.role,
 370            status: message.status.clone(),
 371            timestamp: message.id.0,
 372            cache: message.cache.clone(),
 373        }
 374    }
 375}
 376
 377impl MessageMetadata {
 378    pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range<usize>) -> bool {
 379        let result = match &self.cache {
 380            Some(MessageCacheMetadata { cached_at, .. }) => !buffer.has_edits_since_in_range(
 381                &cached_at,
 382                Range {
 383                    start: buffer.anchor_at(range.start, Bias::Right),
 384                    end: buffer.anchor_at(range.end, Bias::Left),
 385                },
 386            ),
 387            _ => false,
 388        };
 389        result
 390    }
 391}
 392
 393#[derive(Clone, Debug)]
 394pub struct Message {
 395    pub offset_range: Range<usize>,
 396    pub index_range: Range<usize>,
 397    pub anchor_range: Range<language::Anchor>,
 398    pub id: MessageId,
 399    pub role: Role,
 400    pub status: MessageStatus,
 401    pub cache: Option<MessageCacheMetadata>,
 402}
 403
 404#[derive(Debug, Clone)]
 405pub enum Content {
 406    Image {
 407        anchor: language::Anchor,
 408        image_id: u64,
 409        render_image: Arc<RenderImage>,
 410        image: Shared<Task<Option<LanguageModelImage>>>,
 411    },
 412    ToolUse {
 413        range: Range<language::Anchor>,
 414        tool_use: LanguageModelToolUse,
 415    },
 416    ToolResult {
 417        range: Range<language::Anchor>,
 418        tool_use_id: Arc<str>,
 419    },
 420}
 421
 422impl Content {
 423    fn range(&self) -> Range<language::Anchor> {
 424        match self {
 425            Self::Image { anchor, .. } => *anchor..*anchor,
 426            Self::ToolUse { range, .. } | Self::ToolResult { range, .. } => range.clone(),
 427        }
 428    }
 429
 430    fn cmp(&self, other: &Self, buffer: &BufferSnapshot) -> Ordering {
 431        let self_range = self.range();
 432        let other_range = other.range();
 433        if self_range.end.cmp(&other_range.start, buffer).is_lt() {
 434            Ordering::Less
 435        } else if self_range.start.cmp(&other_range.end, buffer).is_gt() {
 436            Ordering::Greater
 437        } else {
 438            Ordering::Equal
 439        }
 440    }
 441}
 442
 443struct PendingCompletion {
 444    id: usize,
 445    assistant_message_id: MessageId,
 446    _task: Task<()>,
 447}
 448
 449#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
 450pub struct SlashCommandId(clock::Lamport);
 451
 452#[derive(Clone, Debug)]
 453pub struct XmlTag {
 454    pub kind: XmlTagKind,
 455    pub range: Range<text::Anchor>,
 456    pub is_open_tag: bool,
 457}
 458
 459#[derive(Copy, Clone, Debug, strum::EnumString, PartialEq, Eq, strum::AsRefStr)]
 460#[strum(serialize_all = "snake_case")]
 461pub enum XmlTagKind {
 462    Patch,
 463    Title,
 464    Edit,
 465    Path,
 466    Description,
 467    OldText,
 468    NewText,
 469    Operation,
 470}
 471
 472pub struct Context {
 473    id: ContextId,
 474    timestamp: clock::Lamport,
 475    version: clock::Global,
 476    pending_ops: Vec<ContextOperation>,
 477    operations: Vec<ContextOperation>,
 478    buffer: Model<Buffer>,
 479    pending_slash_commands: Vec<PendingSlashCommand>,
 480    edits_since_last_parse: language::Subscription,
 481    finished_slash_commands: HashSet<SlashCommandId>,
 482    slash_command_output_sections: Vec<SlashCommandOutputSection<language::Anchor>>,
 483    pending_tool_uses_by_id: HashMap<Arc<str>, PendingToolUse>,
 484    message_anchors: Vec<MessageAnchor>,
 485    contents: Vec<Content>,
 486    messages_metadata: HashMap<MessageId, MessageMetadata>,
 487    summary: Option<ContextSummary>,
 488    pending_summary: Task<Option<()>>,
 489    completion_count: usize,
 490    pending_completions: Vec<PendingCompletion>,
 491    token_count: Option<usize>,
 492    pending_token_count: Task<Option<()>>,
 493    pending_save: Task<Result<()>>,
 494    pending_cache_warming_task: Task<Option<()>>,
 495    path: Option<PathBuf>,
 496    _subscriptions: Vec<Subscription>,
 497    telemetry: Option<Arc<Telemetry>>,
 498    language_registry: Arc<LanguageRegistry>,
 499    patches: Vec<AssistantPatch>,
 500    xml_tags: Vec<XmlTag>,
 501    project: Option<Model<Project>>,
 502    prompt_builder: Arc<PromptBuilder>,
 503}
 504
 505trait ContextAnnotation {
 506    fn range(&self) -> &Range<language::Anchor>;
 507}
 508
 509impl ContextAnnotation for PendingSlashCommand {
 510    fn range(&self) -> &Range<language::Anchor> {
 511        &self.source_range
 512    }
 513}
 514
 515impl ContextAnnotation for AssistantPatch {
 516    fn range(&self) -> &Range<language::Anchor> {
 517        &self.range
 518    }
 519}
 520
 521impl ContextAnnotation for XmlTag {
 522    fn range(&self) -> &Range<language::Anchor> {
 523        &self.range
 524    }
 525}
 526
 527impl EventEmitter<ContextEvent> for Context {}
 528
 529impl Context {
 530    pub fn local(
 531        language_registry: Arc<LanguageRegistry>,
 532        project: Option<Model<Project>>,
 533        telemetry: Option<Arc<Telemetry>>,
 534        prompt_builder: Arc<PromptBuilder>,
 535        cx: &mut ModelContext<Self>,
 536    ) -> Self {
 537        Self::new(
 538            ContextId::new(),
 539            ReplicaId::default(),
 540            language::Capability::ReadWrite,
 541            language_registry,
 542            prompt_builder,
 543            project,
 544            telemetry,
 545            cx,
 546        )
 547    }
 548
 549    #[allow(clippy::too_many_arguments)]
 550    pub fn new(
 551        id: ContextId,
 552        replica_id: ReplicaId,
 553        capability: language::Capability,
 554        language_registry: Arc<LanguageRegistry>,
 555        prompt_builder: Arc<PromptBuilder>,
 556        project: Option<Model<Project>>,
 557        telemetry: Option<Arc<Telemetry>>,
 558        cx: &mut ModelContext<Self>,
 559    ) -> Self {
 560        let buffer = cx.new_model(|_cx| {
 561            let buffer = Buffer::remote(
 562                language::BufferId::new(1).unwrap(),
 563                replica_id,
 564                capability,
 565                "",
 566            );
 567            buffer.set_language_registry(language_registry.clone());
 568            buffer
 569        });
 570        let edits_since_last_slash_command_parse =
 571            buffer.update(cx, |buffer, _| buffer.subscribe());
 572        let mut this = Self {
 573            id,
 574            timestamp: clock::Lamport::new(replica_id),
 575            version: clock::Global::new(),
 576            pending_ops: Vec::new(),
 577            operations: Vec::new(),
 578            message_anchors: Default::default(),
 579            contents: Default::default(),
 580            messages_metadata: Default::default(),
 581            pending_slash_commands: Vec::new(),
 582            finished_slash_commands: HashSet::default(),
 583            pending_tool_uses_by_id: HashMap::default(),
 584            slash_command_output_sections: Vec::new(),
 585            edits_since_last_parse: edits_since_last_slash_command_parse,
 586            summary: None,
 587            pending_summary: Task::ready(None),
 588            completion_count: Default::default(),
 589            pending_completions: Default::default(),
 590            token_count: None,
 591            pending_token_count: Task::ready(None),
 592            pending_cache_warming_task: Task::ready(None),
 593            _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
 594            pending_save: Task::ready(Ok(())),
 595            path: None,
 596            buffer,
 597            telemetry,
 598            project,
 599            language_registry,
 600            patches: Vec::new(),
 601            xml_tags: Vec::new(),
 602            prompt_builder,
 603        };
 604
 605        let first_message_id = MessageId(clock::Lamport {
 606            replica_id: 0,
 607            value: 0,
 608        });
 609        let message = MessageAnchor {
 610            id: first_message_id,
 611            start: language::Anchor::MIN,
 612        };
 613        this.messages_metadata.insert(
 614            first_message_id,
 615            MessageMetadata {
 616                role: Role::User,
 617                status: MessageStatus::Done,
 618                timestamp: first_message_id.0,
 619                cache: None,
 620            },
 621        );
 622        this.message_anchors.push(message);
 623
 624        this.set_language(cx);
 625        this.count_remaining_tokens(cx);
 626        this
 627    }
 628
 629    pub(crate) fn serialize(&self, cx: &AppContext) -> SavedContext {
 630        let buffer = self.buffer.read(cx);
 631        SavedContext {
 632            id: Some(self.id.clone()),
 633            zed: "context".into(),
 634            version: SavedContext::VERSION.into(),
 635            text: buffer.text(),
 636            messages: self
 637                .messages(cx)
 638                .map(|message| SavedMessage {
 639                    id: message.id,
 640                    start: message.offset_range.start,
 641                    metadata: self.messages_metadata[&message.id].clone(),
 642                })
 643                .collect(),
 644            summary: self
 645                .summary
 646                .as_ref()
 647                .map(|summary| summary.text.clone())
 648                .unwrap_or_default(),
 649            slash_command_output_sections: self
 650                .slash_command_output_sections
 651                .iter()
 652                .filter_map(|section| {
 653                    if section.is_valid(buffer) {
 654                        let range = section.range.to_offset(buffer);
 655                        Some(assistant_slash_command::SlashCommandOutputSection {
 656                            range,
 657                            icon: section.icon,
 658                            label: section.label.clone(),
 659                            metadata: section.metadata.clone(),
 660                        })
 661                    } else {
 662                        None
 663                    }
 664                })
 665                .collect(),
 666        }
 667    }
 668
 669    #[allow(clippy::too_many_arguments)]
 670    pub fn deserialize(
 671        saved_context: SavedContext,
 672        path: PathBuf,
 673        language_registry: Arc<LanguageRegistry>,
 674        prompt_builder: Arc<PromptBuilder>,
 675        project: Option<Model<Project>>,
 676        telemetry: Option<Arc<Telemetry>>,
 677        cx: &mut ModelContext<Self>,
 678    ) -> Self {
 679        let id = saved_context.id.clone().unwrap_or_else(ContextId::new);
 680        let mut this = Self::new(
 681            id,
 682            ReplicaId::default(),
 683            language::Capability::ReadWrite,
 684            language_registry,
 685            prompt_builder,
 686            project,
 687            telemetry,
 688            cx,
 689        );
 690        this.path = Some(path);
 691        this.buffer.update(cx, |buffer, cx| {
 692            buffer.set_text(saved_context.text.as_str(), cx)
 693        });
 694        let operations = saved_context.into_ops(&this.buffer, cx);
 695        this.apply_ops(operations, cx);
 696        this
 697    }
 698
 699    pub fn id(&self) -> &ContextId {
 700        &self.id
 701    }
 702
 703    pub fn replica_id(&self) -> ReplicaId {
 704        self.timestamp.replica_id
 705    }
 706
 707    pub fn version(&self, cx: &AppContext) -> ContextVersion {
 708        ContextVersion {
 709            context: self.version.clone(),
 710            buffer: self.buffer.read(cx).version(),
 711        }
 712    }
 713
 714    pub fn set_capability(
 715        &mut self,
 716        capability: language::Capability,
 717        cx: &mut ModelContext<Self>,
 718    ) {
 719        self.buffer
 720            .update(cx, |buffer, cx| buffer.set_capability(capability, cx));
 721    }
 722
 723    fn next_timestamp(&mut self) -> clock::Lamport {
 724        let timestamp = self.timestamp.tick();
 725        self.version.observe(timestamp);
 726        timestamp
 727    }
 728
 729    pub fn serialize_ops(
 730        &self,
 731        since: &ContextVersion,
 732        cx: &AppContext,
 733    ) -> Task<Vec<proto::ContextOperation>> {
 734        let buffer_ops = self
 735            .buffer
 736            .read(cx)
 737            .serialize_ops(Some(since.buffer.clone()), cx);
 738
 739        let mut context_ops = self
 740            .operations
 741            .iter()
 742            .filter(|op| !since.context.observed(op.timestamp()))
 743            .cloned()
 744            .collect::<Vec<_>>();
 745        context_ops.extend(self.pending_ops.iter().cloned());
 746
 747        cx.background_executor().spawn(async move {
 748            let buffer_ops = buffer_ops.await;
 749            context_ops.sort_unstable_by_key(|op| op.timestamp());
 750            buffer_ops
 751                .into_iter()
 752                .map(|op| proto::ContextOperation {
 753                    variant: Some(proto::context_operation::Variant::BufferOperation(
 754                        proto::context_operation::BufferOperation {
 755                            operation: Some(op),
 756                        },
 757                    )),
 758                })
 759                .chain(context_ops.into_iter().map(|op| op.to_proto()))
 760                .collect()
 761        })
 762    }
 763
 764    pub fn apply_ops(
 765        &mut self,
 766        ops: impl IntoIterator<Item = ContextOperation>,
 767        cx: &mut ModelContext<Self>,
 768    ) {
 769        let mut buffer_ops = Vec::new();
 770        for op in ops {
 771            match op {
 772                ContextOperation::BufferOperation(buffer_op) => buffer_ops.push(buffer_op),
 773                op @ _ => self.pending_ops.push(op),
 774            }
 775        }
 776        self.buffer
 777            .update(cx, |buffer, cx| buffer.apply_ops(buffer_ops, cx));
 778        self.flush_ops(cx);
 779    }
 780
 781    fn flush_ops(&mut self, cx: &mut ModelContext<Context>) {
 782        let mut changed_messages = HashSet::default();
 783        let mut summary_changed = false;
 784
 785        self.pending_ops.sort_unstable_by_key(|op| op.timestamp());
 786        for op in mem::take(&mut self.pending_ops) {
 787            if !self.can_apply_op(&op, cx) {
 788                self.pending_ops.push(op);
 789                continue;
 790            }
 791
 792            let timestamp = op.timestamp();
 793            match op.clone() {
 794                ContextOperation::InsertMessage {
 795                    anchor, metadata, ..
 796                } => {
 797                    if self.messages_metadata.contains_key(&anchor.id) {
 798                        // We already applied this operation.
 799                    } else {
 800                        changed_messages.insert(anchor.id);
 801                        self.insert_message(anchor, metadata, cx);
 802                    }
 803                }
 804                ContextOperation::UpdateMessage {
 805                    message_id,
 806                    metadata: new_metadata,
 807                    ..
 808                } => {
 809                    let metadata = self.messages_metadata.get_mut(&message_id).unwrap();
 810                    if new_metadata.timestamp > metadata.timestamp {
 811                        *metadata = new_metadata;
 812                        changed_messages.insert(message_id);
 813                    }
 814                }
 815                ContextOperation::UpdateSummary {
 816                    summary: new_summary,
 817                    ..
 818                } => {
 819                    if self
 820                        .summary
 821                        .as_ref()
 822                        .map_or(true, |summary| new_summary.timestamp > summary.timestamp)
 823                    {
 824                        self.summary = Some(new_summary);
 825                        summary_changed = true;
 826                    }
 827                }
 828                ContextOperation::SlashCommandFinished {
 829                    id,
 830                    output_range,
 831                    sections,
 832                    ..
 833                } => {
 834                    if self.finished_slash_commands.insert(id) {
 835                        let buffer = self.buffer.read(cx);
 836                        self.slash_command_output_sections
 837                            .extend(sections.iter().cloned());
 838                        self.slash_command_output_sections
 839                            .sort_by(|a, b| a.range.cmp(&b.range, buffer));
 840                        cx.emit(ContextEvent::SlashCommandFinished {
 841                            output_range,
 842                            sections,
 843                            expand_result: false,
 844                            run_commands_in_output: false,
 845                        });
 846                    }
 847                }
 848                ContextOperation::BufferOperation(_) => unreachable!(),
 849            }
 850
 851            self.version.observe(timestamp);
 852            self.timestamp.observe(timestamp);
 853            self.operations.push(op);
 854        }
 855
 856        if !changed_messages.is_empty() {
 857            self.message_roles_updated(changed_messages, cx);
 858            cx.emit(ContextEvent::MessagesEdited);
 859            cx.notify();
 860        }
 861
 862        if summary_changed {
 863            cx.emit(ContextEvent::SummaryChanged);
 864            cx.notify();
 865        }
 866    }
 867
 868    fn can_apply_op(&self, op: &ContextOperation, cx: &AppContext) -> bool {
 869        if !self.version.observed_all(op.version()) {
 870            return false;
 871        }
 872
 873        match op {
 874            ContextOperation::InsertMessage { anchor, .. } => self
 875                .buffer
 876                .read(cx)
 877                .version
 878                .observed(anchor.start.timestamp),
 879            ContextOperation::UpdateMessage { message_id, .. } => {
 880                self.messages_metadata.contains_key(message_id)
 881            }
 882            ContextOperation::UpdateSummary { .. } => true,
 883            ContextOperation::SlashCommandFinished {
 884                output_range,
 885                sections,
 886                ..
 887            } => {
 888                let version = &self.buffer.read(cx).version;
 889                sections
 890                    .iter()
 891                    .map(|section| &section.range)
 892                    .chain([output_range])
 893                    .all(|range| {
 894                        let observed_start = range.start == language::Anchor::MIN
 895                            || range.start == language::Anchor::MAX
 896                            || version.observed(range.start.timestamp);
 897                        let observed_end = range.end == language::Anchor::MIN
 898                            || range.end == language::Anchor::MAX
 899                            || version.observed(range.end.timestamp);
 900                        observed_start && observed_end
 901                    })
 902            }
 903            ContextOperation::BufferOperation(_) => {
 904                panic!("buffer operations should always be applied")
 905            }
 906        }
 907    }
 908
 909    fn push_op(&mut self, op: ContextOperation, cx: &mut ModelContext<Self>) {
 910        self.operations.push(op.clone());
 911        cx.emit(ContextEvent::Operation(op));
 912    }
 913
 914    pub fn buffer(&self) -> &Model<Buffer> {
 915        &self.buffer
 916    }
 917
 918    pub fn language_registry(&self) -> Arc<LanguageRegistry> {
 919        self.language_registry.clone()
 920    }
 921
 922    pub fn project(&self) -> Option<Model<Project>> {
 923        self.project.clone()
 924    }
 925
 926    pub fn prompt_builder(&self) -> Arc<PromptBuilder> {
 927        self.prompt_builder.clone()
 928    }
 929
 930    pub fn path(&self) -> Option<&Path> {
 931        self.path.as_deref()
 932    }
 933
 934    pub fn summary(&self) -> Option<&ContextSummary> {
 935        self.summary.as_ref()
 936    }
 937
 938    pub(crate) fn patch_containing(
 939        &self,
 940        position: Point,
 941        cx: &AppContext,
 942    ) -> Option<&AssistantPatch> {
 943        let buffer = self.buffer.read(cx);
 944        let index = self.patches.binary_search_by(|patch| {
 945            let patch_range = patch.range.to_point(&buffer);
 946            if position < patch_range.start {
 947                Ordering::Greater
 948            } else if position > patch_range.end {
 949                Ordering::Less
 950            } else {
 951                Ordering::Equal
 952            }
 953        });
 954        if let Ok(ix) = index {
 955            Some(&self.patches[ix])
 956        } else {
 957            None
 958        }
 959    }
 960
 961    pub fn patch_ranges(&self) -> impl Iterator<Item = Range<language::Anchor>> + '_ {
 962        self.patches.iter().map(|patch| patch.range.clone())
 963    }
 964
 965    pub(crate) fn patch_for_range(
 966        &self,
 967        range: &Range<language::Anchor>,
 968        cx: &AppContext,
 969    ) -> Option<&AssistantPatch> {
 970        let buffer = self.buffer.read(cx);
 971        let index = self.patch_index_for_range(range, buffer).ok()?;
 972        Some(&self.patches[index])
 973    }
 974
 975    fn patch_index_for_range(
 976        &self,
 977        tagged_range: &Range<text::Anchor>,
 978        buffer: &text::BufferSnapshot,
 979    ) -> Result<usize, usize> {
 980        self.patches
 981            .binary_search_by(|probe| probe.range.cmp(&tagged_range, buffer))
 982    }
 983
 984    pub fn pending_slash_commands(&self) -> &[PendingSlashCommand] {
 985        &self.pending_slash_commands
 986    }
 987
 988    pub fn slash_command_output_sections(&self) -> &[SlashCommandOutputSection<language::Anchor>] {
 989        &self.slash_command_output_sections
 990    }
 991
 992    pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {
 993        self.pending_tool_uses_by_id.values().collect()
 994    }
 995
 996    pub fn get_tool_use_by_id(&self, id: &Arc<str>) -> Option<&PendingToolUse> {
 997        self.pending_tool_uses_by_id.get(id)
 998    }
 999
1000    fn set_language(&mut self, cx: &mut ModelContext<Self>) {
1001        let markdown = self.language_registry.language_for_name("Markdown");
1002        cx.spawn(|this, mut cx| async move {
1003            let markdown = markdown.await?;
1004            this.update(&mut cx, |this, cx| {
1005                this.buffer
1006                    .update(cx, |buffer, cx| buffer.set_language(Some(markdown), cx));
1007            })
1008        })
1009        .detach_and_log_err(cx);
1010    }
1011
1012    fn handle_buffer_event(
1013        &mut self,
1014        _: Model<Buffer>,
1015        event: &language::BufferEvent,
1016        cx: &mut ModelContext<Self>,
1017    ) {
1018        match event {
1019            language::BufferEvent::Operation {
1020                operation,
1021                is_local: true,
1022            } => cx.emit(ContextEvent::Operation(ContextOperation::BufferOperation(
1023                operation.clone(),
1024            ))),
1025            language::BufferEvent::Edited => {
1026                self.count_remaining_tokens(cx);
1027                self.reparse(cx);
1028                cx.emit(ContextEvent::MessagesEdited);
1029            }
1030            _ => {}
1031        }
1032    }
1033
1034    pub(crate) fn token_count(&self) -> Option<usize> {
1035        self.token_count
1036    }
1037
1038    pub(crate) fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
1039        let request = self.to_completion_request(RequestType::SuggestEdits, cx); // Conservatively assume SuggestEdits, since it takes more tokens.
1040        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1041            return;
1042        };
1043        self.pending_token_count = cx.spawn(|this, mut cx| {
1044            async move {
1045                cx.background_executor()
1046                    .timer(Duration::from_millis(200))
1047                    .await;
1048
1049                let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
1050                this.update(&mut cx, |this, cx| {
1051                    this.token_count = Some(token_count);
1052                    this.start_cache_warming(&model, cx);
1053                    cx.notify()
1054                })
1055            }
1056            .log_err()
1057        });
1058    }
1059
1060    pub fn mark_cache_anchors(
1061        &mut self,
1062        cache_configuration: &Option<LanguageModelCacheConfiguration>,
1063        speculative: bool,
1064        cx: &mut ModelContext<Self>,
1065    ) -> bool {
1066        let cache_configuration =
1067            cache_configuration
1068                .as_ref()
1069                .unwrap_or(&LanguageModelCacheConfiguration {
1070                    max_cache_anchors: 0,
1071                    should_speculate: false,
1072                    min_total_token: 0,
1073                });
1074
1075        let messages: Vec<Message> = self.messages(cx).collect();
1076
1077        let mut sorted_messages = messages.clone();
1078        if speculative {
1079            // Avoid caching the last message if this is a speculative cache fetch as
1080            // it's likely to change.
1081            sorted_messages.pop();
1082        }
1083        sorted_messages.retain(|m| m.role == Role::User);
1084        sorted_messages.sort_by(|a, b| b.offset_range.len().cmp(&a.offset_range.len()));
1085
1086        let cache_anchors = if self.token_count.unwrap_or(0) < cache_configuration.min_total_token {
1087            // If we have't hit the minimum threshold to enable caching, don't cache anything.
1088            0
1089        } else {
1090            // Save 1 anchor for the inline assistant to use.
1091            max(cache_configuration.max_cache_anchors, 1) - 1
1092        };
1093        sorted_messages.truncate(cache_anchors);
1094
1095        let anchors: HashSet<MessageId> = sorted_messages
1096            .into_iter()
1097            .map(|message| message.id)
1098            .collect();
1099
1100        let buffer = self.buffer.read(cx).snapshot();
1101        let invalidated_caches: HashSet<MessageId> = messages
1102            .iter()
1103            .scan(false, |encountered_invalid, message| {
1104                let message_id = message.id;
1105                let is_invalid = self
1106                    .messages_metadata
1107                    .get(&message_id)
1108                    .map_or(true, |metadata| {
1109                        !metadata.is_cache_valid(&buffer, &message.offset_range)
1110                            || *encountered_invalid
1111                    });
1112                *encountered_invalid |= is_invalid;
1113                Some(if is_invalid { Some(message_id) } else { None })
1114            })
1115            .flatten()
1116            .collect();
1117
1118        let last_anchor = messages.iter().rev().find_map(|message| {
1119            if anchors.contains(&message.id) {
1120                Some(message.id)
1121            } else {
1122                None
1123            }
1124        });
1125
1126        let mut new_anchor_needs_caching = false;
1127        let current_version = &buffer.version;
1128        // If we have no anchors, mark all messages as not being cached.
1129        let mut hit_last_anchor = last_anchor.is_none();
1130
1131        for message in messages.iter() {
1132            if hit_last_anchor {
1133                self.update_metadata(message.id, cx, |metadata| metadata.cache = None);
1134                continue;
1135            }
1136
1137            if let Some(last_anchor) = last_anchor {
1138                if message.id == last_anchor {
1139                    hit_last_anchor = true;
1140                }
1141            }
1142
1143            new_anchor_needs_caching = new_anchor_needs_caching
1144                || (invalidated_caches.contains(&message.id) && anchors.contains(&message.id));
1145
1146            self.update_metadata(message.id, cx, |metadata| {
1147                let cache_status = if invalidated_caches.contains(&message.id) {
1148                    CacheStatus::Pending
1149                } else {
1150                    metadata
1151                        .cache
1152                        .as_ref()
1153                        .map_or(CacheStatus::Pending, |cm| cm.status.clone())
1154                };
1155                metadata.cache = Some(MessageCacheMetadata {
1156                    is_anchor: anchors.contains(&message.id),
1157                    is_final_anchor: hit_last_anchor,
1158                    status: cache_status,
1159                    cached_at: current_version.clone(),
1160                });
1161            });
1162        }
1163        new_anchor_needs_caching
1164    }
1165
1166    fn start_cache_warming(&mut self, model: &Arc<dyn LanguageModel>, cx: &mut ModelContext<Self>) {
1167        let cache_configuration = model.cache_configuration();
1168
1169        if !self.mark_cache_anchors(&cache_configuration, true, cx) {
1170            return;
1171        }
1172        if !self.pending_completions.is_empty() {
1173            return;
1174        }
1175        if let Some(cache_configuration) = cache_configuration {
1176            if !cache_configuration.should_speculate {
1177                return;
1178            }
1179        }
1180
1181        let request = {
1182            let mut req = self.to_completion_request(RequestType::Chat, cx);
1183            // Skip the last message because it's likely to change and
1184            // therefore would be a waste to cache.
1185            req.messages.pop();
1186            req.messages.push(LanguageModelRequestMessage {
1187                role: Role::User,
1188                content: vec!["Respond only with OK, nothing else.".into()],
1189                cache: false,
1190            });
1191            req
1192        };
1193
1194        let model = Arc::clone(model);
1195        self.pending_cache_warming_task = cx.spawn(|this, mut cx| {
1196            async move {
1197                match model.stream_completion(request, &cx).await {
1198                    Ok(mut stream) => {
1199                        stream.next().await;
1200                        log::info!("Cache warming completed successfully");
1201                    }
1202                    Err(e) => {
1203                        log::warn!("Cache warming failed: {}", e);
1204                    }
1205                };
1206                this.update(&mut cx, |this, cx| {
1207                    this.update_cache_status_for_completion(cx);
1208                })
1209                .ok();
1210                anyhow::Ok(())
1211            }
1212            .log_err()
1213        });
1214    }
1215
1216    pub fn update_cache_status_for_completion(&mut self, cx: &mut ModelContext<Self>) {
1217        let cached_message_ids: Vec<MessageId> = self
1218            .messages_metadata
1219            .iter()
1220            .filter_map(|(message_id, metadata)| {
1221                metadata.cache.as_ref().and_then(|cache| {
1222                    if cache.status == CacheStatus::Pending {
1223                        Some(*message_id)
1224                    } else {
1225                        None
1226                    }
1227                })
1228            })
1229            .collect();
1230
1231        for message_id in cached_message_ids {
1232            self.update_metadata(message_id, cx, |metadata| {
1233                if let Some(cache) = &mut metadata.cache {
1234                    cache.status = CacheStatus::Cached;
1235                }
1236            });
1237        }
1238        cx.notify();
1239    }
1240
1241    pub fn reparse(&mut self, cx: &mut ModelContext<Self>) {
1242        let buffer = self.buffer.read(cx).text_snapshot();
1243        let mut row_ranges = self
1244            .edits_since_last_parse
1245            .consume()
1246            .into_iter()
1247            .map(|edit| {
1248                let start_row = buffer.offset_to_point(edit.new.start).row;
1249                let end_row = buffer.offset_to_point(edit.new.end).row + 1;
1250                start_row..end_row
1251            })
1252            .peekable();
1253
1254        let mut removed_slash_command_ranges = Vec::new();
1255        let mut updated_slash_commands = Vec::new();
1256        let mut removed_patches = Vec::new();
1257        let mut updated_patches = Vec::new();
1258        while let Some(mut row_range) = row_ranges.next() {
1259            while let Some(next_row_range) = row_ranges.peek() {
1260                if row_range.end >= next_row_range.start {
1261                    row_range.end = next_row_range.end;
1262                    row_ranges.next();
1263                } else {
1264                    break;
1265                }
1266            }
1267
1268            let start = buffer.anchor_before(Point::new(row_range.start, 0));
1269            let end = buffer.anchor_after(Point::new(
1270                row_range.end - 1,
1271                buffer.line_len(row_range.end - 1),
1272            ));
1273
1274            self.reparse_slash_commands_in_range(
1275                start..end,
1276                &buffer,
1277                &mut updated_slash_commands,
1278                &mut removed_slash_command_ranges,
1279                cx,
1280            );
1281            self.reparse_patches_in_range(
1282                start..end,
1283                &buffer,
1284                &mut updated_patches,
1285                &mut removed_patches,
1286                cx,
1287            );
1288        }
1289
1290        if !updated_slash_commands.is_empty() || !removed_slash_command_ranges.is_empty() {
1291            cx.emit(ContextEvent::PendingSlashCommandsUpdated {
1292                removed: removed_slash_command_ranges,
1293                updated: updated_slash_commands,
1294            });
1295        }
1296
1297        if !updated_patches.is_empty() || !removed_patches.is_empty() {
1298            cx.emit(ContextEvent::PatchesUpdated {
1299                removed: removed_patches,
1300                updated: updated_patches,
1301            });
1302        }
1303    }
1304
1305    fn reparse_slash_commands_in_range(
1306        &mut self,
1307        range: Range<text::Anchor>,
1308        buffer: &BufferSnapshot,
1309        updated: &mut Vec<PendingSlashCommand>,
1310        removed: &mut Vec<Range<text::Anchor>>,
1311        cx: &AppContext,
1312    ) {
1313        let old_range = self.pending_command_indices_for_range(range.clone(), cx);
1314
1315        let mut new_commands = Vec::new();
1316        let mut lines = buffer.text_for_range(range).lines();
1317        let mut offset = lines.offset();
1318        while let Some(line) = lines.next() {
1319            if let Some(command_line) = SlashCommandLine::parse(line) {
1320                let name = &line[command_line.name.clone()];
1321                let arguments = command_line
1322                    .arguments
1323                    .iter()
1324                    .filter_map(|argument_range| {
1325                        if argument_range.is_empty() {
1326                            None
1327                        } else {
1328                            line.get(argument_range.clone())
1329                        }
1330                    })
1331                    .map(ToOwned::to_owned)
1332                    .collect::<SmallVec<_>>();
1333                if let Some(command) = SlashCommandRegistry::global(cx).command(name) {
1334                    if !command.requires_argument() || !arguments.is_empty() {
1335                        let start_ix = offset + command_line.name.start - 1;
1336                        let end_ix = offset
1337                            + command_line
1338                                .arguments
1339                                .last()
1340                                .map_or(command_line.name.end, |argument| argument.end);
1341                        let source_range =
1342                            buffer.anchor_after(start_ix)..buffer.anchor_after(end_ix);
1343                        let pending_command = PendingSlashCommand {
1344                            name: name.to_string(),
1345                            arguments,
1346                            source_range,
1347                            status: PendingSlashCommandStatus::Idle,
1348                        };
1349                        updated.push(pending_command.clone());
1350                        new_commands.push(pending_command);
1351                    }
1352                }
1353            }
1354
1355            offset = lines.offset();
1356        }
1357
1358        let removed_commands = self.pending_slash_commands.splice(old_range, new_commands);
1359        removed.extend(removed_commands.map(|command| command.source_range));
1360    }
1361
1362    fn reparse_patches_in_range(
1363        &mut self,
1364        range: Range<text::Anchor>,
1365        buffer: &BufferSnapshot,
1366        updated: &mut Vec<Range<text::Anchor>>,
1367        removed: &mut Vec<Range<text::Anchor>>,
1368        cx: &mut ModelContext<Self>,
1369    ) {
1370        // Rebuild the XML tags in the edited range.
1371        let intersecting_tags_range =
1372            self.indices_intersecting_buffer_range(&self.xml_tags, range.clone(), cx);
1373        let new_tags = self.parse_xml_tags_in_range(buffer, range.clone(), cx);
1374        self.xml_tags
1375            .splice(intersecting_tags_range.clone(), new_tags);
1376
1377        // Find which patches intersect the changed range.
1378        let intersecting_patches_range =
1379            self.indices_intersecting_buffer_range(&self.patches, range.clone(), cx);
1380
1381        // Reparse all tags after the last unchanged patch before the change.
1382        let mut tags_start_ix = 0;
1383        if let Some(preceding_unchanged_patch) =
1384            self.patches[..intersecting_patches_range.start].last()
1385        {
1386            tags_start_ix = match self.xml_tags.binary_search_by(|tag| {
1387                tag.range
1388                    .start
1389                    .cmp(&preceding_unchanged_patch.range.end, buffer)
1390                    .then(Ordering::Less)
1391            }) {
1392                Ok(ix) | Err(ix) => ix,
1393            };
1394        }
1395
1396        // Rebuild the patches in the range.
1397        let new_patches = self.parse_patches(tags_start_ix, range.end, buffer, cx);
1398        updated.extend(new_patches.iter().map(|patch| patch.range.clone()));
1399        let removed_patches = self.patches.splice(intersecting_patches_range, new_patches);
1400        removed.extend(
1401            removed_patches
1402                .map(|patch| patch.range)
1403                .filter(|range| !updated.contains(&range)),
1404        );
1405    }
1406
1407    fn parse_xml_tags_in_range(
1408        &self,
1409        buffer: &BufferSnapshot,
1410        range: Range<text::Anchor>,
1411        cx: &AppContext,
1412    ) -> Vec<XmlTag> {
1413        let mut messages = self.messages(cx).peekable();
1414
1415        let mut tags = Vec::new();
1416        let mut lines = buffer.text_for_range(range).lines();
1417        let mut offset = lines.offset();
1418
1419        while let Some(line) = lines.next() {
1420            while let Some(message) = messages.peek() {
1421                if offset < message.offset_range.end {
1422                    break;
1423                } else {
1424                    messages.next();
1425                }
1426            }
1427
1428            let is_assistant_message = messages
1429                .peek()
1430                .map_or(false, |message| message.role == Role::Assistant);
1431            if is_assistant_message {
1432                for (start_ix, _) in line.match_indices('<') {
1433                    let mut name_start_ix = start_ix + 1;
1434                    let closing_bracket_ix = line[start_ix..].find('>').map(|i| start_ix + i);
1435                    if let Some(closing_bracket_ix) = closing_bracket_ix {
1436                        let end_ix = closing_bracket_ix + 1;
1437                        let mut is_open_tag = true;
1438                        if line[name_start_ix..closing_bracket_ix].starts_with('/') {
1439                            name_start_ix += 1;
1440                            is_open_tag = false;
1441                        }
1442                        let tag_inner = &line[name_start_ix..closing_bracket_ix];
1443                        let tag_name_len = tag_inner
1444                            .find(|c: char| c.is_whitespace())
1445                            .unwrap_or(tag_inner.len());
1446                        if let Ok(kind) = XmlTagKind::from_str(&tag_inner[..tag_name_len]) {
1447                            tags.push(XmlTag {
1448                                range: buffer.anchor_after(offset + start_ix)
1449                                    ..buffer.anchor_before(offset + end_ix),
1450                                is_open_tag,
1451                                kind,
1452                            });
1453                        };
1454                    }
1455                }
1456            }
1457
1458            offset = lines.offset();
1459        }
1460        tags
1461    }
1462
1463    fn parse_patches(
1464        &mut self,
1465        tags_start_ix: usize,
1466        buffer_end: text::Anchor,
1467        buffer: &BufferSnapshot,
1468        cx: &AppContext,
1469    ) -> Vec<AssistantPatch> {
1470        let mut new_patches = Vec::new();
1471        let mut pending_patch = None;
1472        let mut patch_tag_depth = 0;
1473        let mut tags = self.xml_tags[tags_start_ix..].iter().peekable();
1474        'tags: while let Some(tag) = tags.next() {
1475            if tag.range.start.cmp(&buffer_end, buffer).is_gt() && patch_tag_depth == 0 {
1476                break;
1477            }
1478
1479            if tag.kind == XmlTagKind::Patch && tag.is_open_tag {
1480                patch_tag_depth += 1;
1481                let patch_start = tag.range.start;
1482                let mut edits = Vec::<Result<AssistantEdit>>::new();
1483                let mut patch = AssistantPatch {
1484                    range: patch_start..patch_start,
1485                    title: String::new().into(),
1486                    edits: Default::default(),
1487                    status: crate::AssistantPatchStatus::Pending,
1488                };
1489
1490                while let Some(tag) = tags.next() {
1491                    if tag.kind == XmlTagKind::Patch && !tag.is_open_tag {
1492                        patch_tag_depth -= 1;
1493                        if patch_tag_depth == 0 {
1494                            patch.range.end = tag.range.end;
1495
1496                            // Include the line immediately after this <patch> tag if it's empty.
1497                            let patch_end_offset = patch.range.end.to_offset(buffer);
1498                            let mut patch_end_chars = buffer.chars_at(patch_end_offset);
1499                            if patch_end_chars.next() == Some('\n')
1500                                && patch_end_chars.next().map_or(true, |ch| ch == '\n')
1501                            {
1502                                let messages = self.messages_for_offsets(
1503                                    [patch_end_offset, patch_end_offset + 1],
1504                                    cx,
1505                                );
1506                                if messages.len() == 1 {
1507                                    patch.range.end = buffer.anchor_before(patch_end_offset + 1);
1508                                }
1509                            }
1510
1511                            edits.sort_unstable_by(|a, b| {
1512                                if let (Ok(a), Ok(b)) = (a, b) {
1513                                    a.path.cmp(&b.path)
1514                                } else {
1515                                    Ordering::Equal
1516                                }
1517                            });
1518                            patch.edits = edits.into();
1519                            patch.status = AssistantPatchStatus::Ready;
1520                            new_patches.push(patch);
1521                            continue 'tags;
1522                        }
1523                    }
1524
1525                    if tag.kind == XmlTagKind::Title && tag.is_open_tag {
1526                        let content_start = tag.range.end;
1527                        while let Some(tag) = tags.next() {
1528                            if tag.kind == XmlTagKind::Title && !tag.is_open_tag {
1529                                let content_end = tag.range.start;
1530                                patch.title =
1531                                    trimmed_text_in_range(buffer, content_start..content_end)
1532                                        .into();
1533                                break;
1534                            }
1535                        }
1536                    }
1537
1538                    if tag.kind == XmlTagKind::Edit && tag.is_open_tag {
1539                        let mut path = None;
1540                        let mut old_text = None;
1541                        let mut new_text = None;
1542                        let mut operation = None;
1543                        let mut description = None;
1544
1545                        while let Some(tag) = tags.next() {
1546                            if tag.kind == XmlTagKind::Edit && !tag.is_open_tag {
1547                                edits.push(AssistantEdit::new(
1548                                    path,
1549                                    operation,
1550                                    old_text,
1551                                    new_text,
1552                                    description,
1553                                ));
1554                                break;
1555                            }
1556
1557                            if tag.is_open_tag
1558                                && [
1559                                    XmlTagKind::Path,
1560                                    XmlTagKind::OldText,
1561                                    XmlTagKind::NewText,
1562                                    XmlTagKind::Operation,
1563                                    XmlTagKind::Description,
1564                                ]
1565                                .contains(&tag.kind)
1566                            {
1567                                let kind = tag.kind;
1568                                let content_start = tag.range.end;
1569                                if let Some(tag) = tags.peek() {
1570                                    if tag.kind == kind && !tag.is_open_tag {
1571                                        let tag = tags.next().unwrap();
1572                                        let content_end = tag.range.start;
1573                                        let content = trimmed_text_in_range(
1574                                            buffer,
1575                                            content_start..content_end,
1576                                        );
1577                                        match kind {
1578                                            XmlTagKind::Path => path = Some(content),
1579                                            XmlTagKind::Operation => operation = Some(content),
1580                                            XmlTagKind::OldText => {
1581                                                old_text = Some(content).filter(|s| !s.is_empty())
1582                                            }
1583                                            XmlTagKind::NewText => {
1584                                                new_text = Some(content).filter(|s| !s.is_empty())
1585                                            }
1586                                            XmlTagKind::Description => {
1587                                                description =
1588                                                    Some(content).filter(|s| !s.is_empty())
1589                                            }
1590                                            _ => {}
1591                                        }
1592                                    }
1593                                }
1594                            }
1595                        }
1596                    }
1597                }
1598
1599                patch.edits = edits.into();
1600                pending_patch = Some(patch);
1601            }
1602        }
1603
1604        if let Some(mut pending_patch) = pending_patch {
1605            let patch_start = pending_patch.range.start.to_offset(buffer);
1606            if let Some(message) = self.message_for_offset(patch_start, cx) {
1607                if message.anchor_range.end == text::Anchor::MAX {
1608                    pending_patch.range.end = text::Anchor::MAX;
1609                } else {
1610                    let message_end = buffer.anchor_after(message.offset_range.end - 1);
1611                    pending_patch.range.end = message_end;
1612                }
1613            } else {
1614                pending_patch.range.end = text::Anchor::MAX;
1615            }
1616
1617            new_patches.push(pending_patch);
1618        }
1619
1620        new_patches
1621    }
1622
1623    pub fn pending_command_for_position(
1624        &mut self,
1625        position: language::Anchor,
1626        cx: &mut ModelContext<Self>,
1627    ) -> Option<&mut PendingSlashCommand> {
1628        let buffer = self.buffer.read(cx);
1629        match self
1630            .pending_slash_commands
1631            .binary_search_by(|probe| probe.source_range.end.cmp(&position, buffer))
1632        {
1633            Ok(ix) => Some(&mut self.pending_slash_commands[ix]),
1634            Err(ix) => {
1635                let cmd = self.pending_slash_commands.get_mut(ix)?;
1636                if position.cmp(&cmd.source_range.start, buffer).is_ge()
1637                    && position.cmp(&cmd.source_range.end, buffer).is_le()
1638                {
1639                    Some(cmd)
1640                } else {
1641                    None
1642                }
1643            }
1644        }
1645    }
1646
1647    pub fn pending_commands_for_range(
1648        &self,
1649        range: Range<language::Anchor>,
1650        cx: &AppContext,
1651    ) -> &[PendingSlashCommand] {
1652        let range = self.pending_command_indices_for_range(range, cx);
1653        &self.pending_slash_commands[range]
1654    }
1655
1656    fn pending_command_indices_for_range(
1657        &self,
1658        range: Range<language::Anchor>,
1659        cx: &AppContext,
1660    ) -> Range<usize> {
1661        self.indices_intersecting_buffer_range(&self.pending_slash_commands, range, cx)
1662    }
1663
1664    fn indices_intersecting_buffer_range<T: ContextAnnotation>(
1665        &self,
1666        all_annotations: &[T],
1667        range: Range<language::Anchor>,
1668        cx: &AppContext,
1669    ) -> Range<usize> {
1670        let buffer = self.buffer.read(cx);
1671        let start_ix = match all_annotations
1672            .binary_search_by(|probe| probe.range().end.cmp(&range.start, &buffer))
1673        {
1674            Ok(ix) | Err(ix) => ix,
1675        };
1676        let end_ix = match all_annotations
1677            .binary_search_by(|probe| probe.range().start.cmp(&range.end, &buffer))
1678        {
1679            Ok(ix) => ix + 1,
1680            Err(ix) => ix,
1681        };
1682        start_ix..end_ix
1683    }
1684
1685    pub fn insert_command_output(
1686        &mut self,
1687        command_range: Range<language::Anchor>,
1688        output: Task<SlashCommandResult>,
1689        ensure_trailing_newline: bool,
1690        expand_result: bool,
1691        cx: &mut ModelContext<Self>,
1692    ) {
1693        self.reparse(cx);
1694
1695        let insert_output_task = cx.spawn(|this, mut cx| {
1696            let command_range = command_range.clone();
1697            async move {
1698                let output = output.await;
1699                let output = match output {
1700                    Ok(output) => SlashCommandOutput::from_event_stream(output).await,
1701                    Err(err) => Err(err),
1702                };
1703                this.update(&mut cx, |this, cx| match output {
1704                    Ok(mut output) => {
1705                        output.ensure_valid_section_ranges();
1706
1707                        // Ensure there is a newline after the last section.
1708                        if ensure_trailing_newline {
1709                            let has_newline_after_last_section =
1710                                output.sections.last().map_or(false, |last_section| {
1711                                    output.text[last_section.range.end..].ends_with('\n')
1712                                });
1713                            if !has_newline_after_last_section {
1714                                output.text.push('\n');
1715                            }
1716                        }
1717
1718                        let version = this.version.clone();
1719                        let command_id = SlashCommandId(this.next_timestamp());
1720                        let (operation, event) = this.buffer.update(cx, |buffer, cx| {
1721                            let start = command_range.start.to_offset(buffer);
1722                            let old_end = command_range.end.to_offset(buffer);
1723                            let new_end = start + output.text.len();
1724                            buffer.edit([(start..old_end, output.text)], None, cx);
1725
1726                            let mut sections = output
1727                                .sections
1728                                .into_iter()
1729                                .map(|section| SlashCommandOutputSection {
1730                                    range: buffer.anchor_after(start + section.range.start)
1731                                        ..buffer.anchor_before(start + section.range.end),
1732                                    icon: section.icon,
1733                                    label: section.label,
1734                                    metadata: section.metadata,
1735                                })
1736                                .collect::<Vec<_>>();
1737                            sections.sort_by(|a, b| a.range.cmp(&b.range, buffer));
1738
1739                            this.slash_command_output_sections
1740                                .extend(sections.iter().cloned());
1741                            this.slash_command_output_sections
1742                                .sort_by(|a, b| a.range.cmp(&b.range, buffer));
1743
1744                            let output_range =
1745                                buffer.anchor_after(start)..buffer.anchor_before(new_end);
1746                            this.finished_slash_commands.insert(command_id);
1747
1748                            (
1749                                ContextOperation::SlashCommandFinished {
1750                                    id: command_id,
1751                                    output_range: output_range.clone(),
1752                                    sections: sections.clone(),
1753                                    version,
1754                                },
1755                                ContextEvent::SlashCommandFinished {
1756                                    output_range,
1757                                    sections,
1758                                    run_commands_in_output: output.run_commands_in_text,
1759                                    expand_result,
1760                                },
1761                            )
1762                        });
1763
1764                        this.push_op(operation, cx);
1765                        cx.emit(event);
1766                    }
1767                    Err(error) => {
1768                        if let Some(pending_command) =
1769                            this.pending_command_for_position(command_range.start, cx)
1770                        {
1771                            pending_command.status =
1772                                PendingSlashCommandStatus::Error(error.to_string());
1773                            cx.emit(ContextEvent::PendingSlashCommandsUpdated {
1774                                removed: vec![pending_command.source_range.clone()],
1775                                updated: vec![pending_command.clone()],
1776                            });
1777                        }
1778                    }
1779                })
1780                .ok();
1781            }
1782        });
1783
1784        if let Some(pending_command) = self.pending_command_for_position(command_range.start, cx) {
1785            pending_command.status = PendingSlashCommandStatus::Running {
1786                _task: insert_output_task.shared(),
1787            };
1788            cx.emit(ContextEvent::PendingSlashCommandsUpdated {
1789                removed: vec![pending_command.source_range.clone()],
1790                updated: vec![pending_command.clone()],
1791            });
1792        }
1793    }
1794
1795    pub fn insert_tool_output(
1796        &mut self,
1797        tool_use_id: Arc<str>,
1798        output: Task<Result<String>>,
1799        cx: &mut ModelContext<Self>,
1800    ) {
1801        let insert_output_task = cx.spawn(|this, mut cx| {
1802            let tool_use_id = tool_use_id.clone();
1803            async move {
1804                let output = output.await;
1805                this.update(&mut cx, |this, cx| match output {
1806                    Ok(mut output) => {
1807                        const NEWLINE: char = '\n';
1808
1809                        if !output.ends_with(NEWLINE) {
1810                            output.push(NEWLINE);
1811                        }
1812
1813                        let anchor_range = this.buffer.update(cx, |buffer, cx| {
1814                            let insert_start = buffer.len().to_offset(buffer);
1815                            let insert_end = insert_start;
1816
1817                            let start = insert_start;
1818                            let end = start + output.len() - NEWLINE.len_utf8();
1819
1820                            buffer.edit([(insert_start..insert_end, output)], None, cx);
1821
1822                            let output_range = buffer.anchor_after(start)..buffer.anchor_after(end);
1823
1824                            output_range
1825                        });
1826
1827                        this.insert_content(
1828                            Content::ToolResult {
1829                                range: anchor_range.clone(),
1830                                tool_use_id: tool_use_id.clone(),
1831                            },
1832                            cx,
1833                        );
1834
1835                        cx.emit(ContextEvent::ToolFinished {
1836                            tool_use_id,
1837                            output_range: anchor_range,
1838                        });
1839                    }
1840                    Err(err) => {
1841                        if let Some(tool_use) = this.pending_tool_uses_by_id.get_mut(&tool_use_id) {
1842                            tool_use.status = PendingToolUseStatus::Error(err.to_string());
1843                        }
1844                    }
1845                })
1846                .ok();
1847            }
1848        });
1849
1850        if let Some(tool_use) = self.pending_tool_uses_by_id.get_mut(&tool_use_id) {
1851            tool_use.status = PendingToolUseStatus::Running {
1852                _task: insert_output_task.shared(),
1853            };
1854        }
1855    }
1856
1857    pub fn completion_provider_changed(&mut self, cx: &mut ModelContext<Self>) {
1858        self.count_remaining_tokens(cx);
1859    }
1860
1861    fn get_last_valid_message_id(&self, cx: &ModelContext<Self>) -> Option<MessageId> {
1862        self.message_anchors.iter().rev().find_map(|message| {
1863            message
1864                .start
1865                .is_valid(self.buffer.read(cx))
1866                .then_some(message.id)
1867        })
1868    }
1869
1870    pub fn assist(
1871        &mut self,
1872        request_type: RequestType,
1873        cx: &mut ModelContext<Self>,
1874    ) -> Option<MessageAnchor> {
1875        let model_registry = LanguageModelRegistry::read_global(cx);
1876        let provider = model_registry.active_provider()?;
1877        let model = model_registry.active_model()?;
1878        let last_message_id = self.get_last_valid_message_id(cx)?;
1879
1880        if !provider.is_authenticated(cx) {
1881            log::info!("completion provider has no credentials");
1882            return None;
1883        }
1884        // Compute which messages to cache, including the last one.
1885        self.mark_cache_anchors(&model.cache_configuration(), false, cx);
1886
1887        let mut request = self.to_completion_request(request_type, cx);
1888
1889        if cx.has_flag::<ToolUseFeatureFlag>() {
1890            let tool_registry = ToolRegistry::global(cx);
1891            request.tools = tool_registry
1892                .tools()
1893                .into_iter()
1894                .map(|tool| LanguageModelRequestTool {
1895                    name: tool.name(),
1896                    description: tool.description(),
1897                    input_schema: tool.input_schema(),
1898                })
1899                .collect();
1900        }
1901
1902        let assistant_message = self
1903            .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
1904            .unwrap();
1905
1906        // Queue up the user's next reply.
1907        let user_message = self
1908            .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
1909            .unwrap();
1910
1911        let pending_completion_id = post_inc(&mut self.completion_count);
1912
1913        let task = cx.spawn({
1914            |this, mut cx| async move {
1915                let stream = model.stream_completion(request, &cx);
1916                let assistant_message_id = assistant_message.id;
1917                let mut response_latency = None;
1918                let stream_completion = async {
1919                    let request_start = Instant::now();
1920                    let mut events = stream.await?;
1921                    let mut stop_reason = StopReason::EndTurn;
1922
1923                    while let Some(event) = events.next().await {
1924                        if response_latency.is_none() {
1925                            response_latency = Some(request_start.elapsed());
1926                        }
1927                        let event = event?;
1928
1929                        this.update(&mut cx, |this, cx| {
1930                            let message_ix = this
1931                                .message_anchors
1932                                .iter()
1933                                .position(|message| message.id == assistant_message_id)?;
1934                            this.buffer.update(cx, |buffer, cx| {
1935                                let message_old_end_offset = this.message_anchors[message_ix + 1..]
1936                                    .iter()
1937                                    .find(|message| message.start.is_valid(buffer))
1938                                    .map_or(buffer.len(), |message| {
1939                                        message.start.to_offset(buffer).saturating_sub(1)
1940                                    });
1941
1942                                match event {
1943                                    LanguageModelCompletionEvent::Stop(reason) => {
1944                                        stop_reason = reason;
1945                                    }
1946                                    LanguageModelCompletionEvent::Text(chunk) => {
1947                                        buffer.edit(
1948                                            [(
1949                                                message_old_end_offset..message_old_end_offset,
1950                                                chunk,
1951                                            )],
1952                                            None,
1953                                            cx,
1954                                        );
1955                                    }
1956                                    LanguageModelCompletionEvent::ToolUse(tool_use) => {
1957                                        const NEWLINE: char = '\n';
1958
1959                                        let mut text = String::new();
1960                                        text.push(NEWLINE);
1961                                        text.push_str(
1962                                            &serde_json::to_string_pretty(&tool_use)
1963                                                .expect("failed to serialize tool use to JSON"),
1964                                        );
1965                                        text.push(NEWLINE);
1966                                        let text_len = text.len();
1967
1968                                        buffer.edit(
1969                                            [(
1970                                                message_old_end_offset..message_old_end_offset,
1971                                                text,
1972                                            )],
1973                                            None,
1974                                            cx,
1975                                        );
1976
1977                                        let start_ix = message_old_end_offset + NEWLINE.len_utf8();
1978                                        let end_ix =
1979                                            message_old_end_offset + text_len - NEWLINE.len_utf8();
1980                                        let source_range = buffer.anchor_after(start_ix)
1981                                            ..buffer.anchor_after(end_ix);
1982
1983                                        let tool_use_id: Arc<str> = tool_use.id.into();
1984                                        this.pending_tool_uses_by_id.insert(
1985                                            tool_use_id.clone(),
1986                                            PendingToolUse {
1987                                                id: tool_use_id,
1988                                                name: tool_use.name,
1989                                                input: tool_use.input,
1990                                                status: PendingToolUseStatus::Idle,
1991                                                source_range,
1992                                            },
1993                                        );
1994                                    }
1995                                }
1996                            });
1997
1998                            cx.emit(ContextEvent::StreamedCompletion);
1999
2000                            Some(())
2001                        })?;
2002                        smol::future::yield_now().await;
2003                    }
2004                    this.update(&mut cx, |this, cx| {
2005                        this.pending_completions
2006                            .retain(|completion| completion.id != pending_completion_id);
2007                        this.summarize(false, cx);
2008                        this.update_cache_status_for_completion(cx);
2009                    })?;
2010
2011                    anyhow::Ok(stop_reason)
2012                };
2013
2014                let result = stream_completion.await;
2015
2016                this.update(&mut cx, |this, cx| {
2017                    let error_message = if let Some(error) = result.as_ref().err() {
2018                        if error.is::<PaymentRequiredError>() {
2019                            cx.emit(ContextEvent::ShowPaymentRequiredError);
2020                            this.update_metadata(assistant_message_id, cx, |metadata| {
2021                                metadata.status = MessageStatus::Canceled;
2022                            });
2023                            Some(error.to_string())
2024                        } else if error.is::<MaxMonthlySpendReachedError>() {
2025                            cx.emit(ContextEvent::ShowMaxMonthlySpendReachedError);
2026                            this.update_metadata(assistant_message_id, cx, |metadata| {
2027                                metadata.status = MessageStatus::Canceled;
2028                            });
2029                            Some(error.to_string())
2030                        } else {
2031                            let error_message = error.to_string().trim().to_string();
2032                            cx.emit(ContextEvent::ShowAssistError(SharedString::from(
2033                                error_message.clone(),
2034                            )));
2035                            this.update_metadata(assistant_message_id, cx, |metadata| {
2036                                metadata.status =
2037                                    MessageStatus::Error(SharedString::from(error_message.clone()));
2038                            });
2039                            Some(error_message)
2040                        }
2041                    } else {
2042                        this.update_metadata(assistant_message_id, cx, |metadata| {
2043                            metadata.status = MessageStatus::Done;
2044                        });
2045                        None
2046                    };
2047
2048                    if let Some(telemetry) = this.telemetry.as_ref() {
2049                        let language_name = this
2050                            .buffer
2051                            .read(cx)
2052                            .language()
2053                            .map(|language| language.name());
2054                        telemetry.report_assistant_event(AssistantEvent {
2055                            conversation_id: Some(this.id.0.clone()),
2056                            kind: AssistantKind::Panel,
2057                            phase: AssistantPhase::Response,
2058                            model: model.telemetry_id(),
2059                            model_provider: model.provider_id().to_string(),
2060                            response_latency,
2061                            error_message,
2062                            language_name: language_name.map(|name| name.to_proto()),
2063                        });
2064                    }
2065
2066                    if let Ok(stop_reason) = result {
2067                        match stop_reason {
2068                            StopReason::ToolUse => {
2069                                cx.emit(ContextEvent::UsePendingTools);
2070                            }
2071                            StopReason::EndTurn => {}
2072                            StopReason::MaxTokens => {}
2073                        }
2074                    }
2075                })
2076                .ok();
2077            }
2078        });
2079
2080        self.pending_completions.push(PendingCompletion {
2081            id: pending_completion_id,
2082            assistant_message_id: assistant_message.id,
2083            _task: task,
2084        });
2085
2086        Some(user_message)
2087    }
2088
2089    pub fn to_completion_request(
2090        &self,
2091        request_type: RequestType,
2092        cx: &AppContext,
2093    ) -> LanguageModelRequest {
2094        let buffer = self.buffer.read(cx);
2095
2096        let mut contents = self.contents(cx).peekable();
2097
2098        fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<String> {
2099            let text: String = buffer.text_for_range(range.clone()).collect();
2100            if text.trim().is_empty() {
2101                None
2102            } else {
2103                Some(text)
2104            }
2105        }
2106
2107        let mut completion_request = LanguageModelRequest {
2108            messages: Vec::new(),
2109            tools: Vec::new(),
2110            stop: Vec::new(),
2111            temperature: None,
2112        };
2113        for message in self.messages(cx) {
2114            if message.status != MessageStatus::Done {
2115                continue;
2116            }
2117
2118            let mut offset = message.offset_range.start;
2119            let mut request_message = LanguageModelRequestMessage {
2120                role: message.role,
2121                content: Vec::new(),
2122                cache: message
2123                    .cache
2124                    .as_ref()
2125                    .map_or(false, |cache| cache.is_anchor),
2126            };
2127
2128            while let Some(content) = contents.peek() {
2129                if content
2130                    .range()
2131                    .end
2132                    .cmp(&message.anchor_range.end, buffer)
2133                    .is_lt()
2134                {
2135                    let content = contents.next().unwrap();
2136                    let range = content.range().to_offset(buffer);
2137                    request_message.content.extend(
2138                        collect_text_content(buffer, offset..range.start).map(MessageContent::Text),
2139                    );
2140
2141                    match content {
2142                        Content::Image { image, .. } => {
2143                            if let Some(image) = image.clone().now_or_never().flatten() {
2144                                request_message
2145                                    .content
2146                                    .push(language_model::MessageContent::Image(image));
2147                            }
2148                        }
2149                        Content::ToolUse { tool_use, .. } => {
2150                            request_message
2151                                .content
2152                                .push(language_model::MessageContent::ToolUse(tool_use.clone()));
2153                        }
2154                        Content::ToolResult { tool_use_id, .. } => {
2155                            request_message.content.push(
2156                                language_model::MessageContent::ToolResult(
2157                                    LanguageModelToolResult {
2158                                        tool_use_id: tool_use_id.to_string(),
2159                                        is_error: false,
2160                                        content: collect_text_content(buffer, range.clone())
2161                                            .unwrap_or_default(),
2162                                    },
2163                                ),
2164                            );
2165                        }
2166                    }
2167
2168                    offset = range.end;
2169                } else {
2170                    break;
2171                }
2172            }
2173
2174            request_message.content.extend(
2175                collect_text_content(buffer, offset..message.offset_range.end)
2176                    .map(MessageContent::Text),
2177            );
2178
2179            completion_request.messages.push(request_message);
2180        }
2181
2182        if let RequestType::SuggestEdits = request_type {
2183            if let Ok(preamble) = self.prompt_builder.generate_workflow_prompt() {
2184                let last_elem_index = completion_request.messages.len();
2185
2186                completion_request
2187                    .messages
2188                    .push(LanguageModelRequestMessage {
2189                        role: Role::User,
2190                        content: vec![MessageContent::Text(preamble)],
2191                        cache: false,
2192                    });
2193
2194                // The preamble message should be sent right before the last actual user message.
2195                completion_request
2196                    .messages
2197                    .swap(last_elem_index, last_elem_index.saturating_sub(1));
2198            }
2199        }
2200
2201        completion_request
2202    }
2203
2204    pub fn cancel_last_assist(&mut self, cx: &mut ModelContext<Self>) -> bool {
2205        if let Some(pending_completion) = self.pending_completions.pop() {
2206            self.update_metadata(pending_completion.assistant_message_id, cx, |metadata| {
2207                if metadata.status == MessageStatus::Pending {
2208                    metadata.status = MessageStatus::Canceled;
2209                }
2210            });
2211            true
2212        } else {
2213            false
2214        }
2215    }
2216
2217    pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
2218        for id in &ids {
2219            if let Some(metadata) = self.messages_metadata.get(id) {
2220                let role = metadata.role.cycle();
2221                self.update_metadata(*id, cx, |metadata| metadata.role = role);
2222            }
2223        }
2224
2225        self.message_roles_updated(ids, cx);
2226    }
2227
2228    fn message_roles_updated(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
2229        let mut ranges = Vec::new();
2230        for message in self.messages(cx) {
2231            if ids.contains(&message.id) {
2232                ranges.push(message.anchor_range.clone());
2233            }
2234        }
2235
2236        let buffer = self.buffer.read(cx).text_snapshot();
2237        let mut updated = Vec::new();
2238        let mut removed = Vec::new();
2239        for range in ranges {
2240            self.reparse_patches_in_range(range, &buffer, &mut updated, &mut removed, cx);
2241        }
2242
2243        if !updated.is_empty() || !removed.is_empty() {
2244            cx.emit(ContextEvent::PatchesUpdated { removed, updated })
2245        }
2246    }
2247
2248    pub fn update_metadata(
2249        &mut self,
2250        id: MessageId,
2251        cx: &mut ModelContext<Self>,
2252        f: impl FnOnce(&mut MessageMetadata),
2253    ) {
2254        let version = self.version.clone();
2255        let timestamp = self.next_timestamp();
2256        if let Some(metadata) = self.messages_metadata.get_mut(&id) {
2257            f(metadata);
2258            metadata.timestamp = timestamp;
2259            let operation = ContextOperation::UpdateMessage {
2260                message_id: id,
2261                metadata: metadata.clone(),
2262                version,
2263            };
2264            self.push_op(operation, cx);
2265            cx.emit(ContextEvent::MessagesEdited);
2266            cx.notify();
2267        }
2268    }
2269
2270    pub fn insert_message_after(
2271        &mut self,
2272        message_id: MessageId,
2273        role: Role,
2274        status: MessageStatus,
2275        cx: &mut ModelContext<Self>,
2276    ) -> Option<MessageAnchor> {
2277        if let Some(prev_message_ix) = self
2278            .message_anchors
2279            .iter()
2280            .position(|message| message.id == message_id)
2281        {
2282            // Find the next valid message after the one we were given.
2283            let mut next_message_ix = prev_message_ix + 1;
2284            while let Some(next_message) = self.message_anchors.get(next_message_ix) {
2285                if next_message.start.is_valid(self.buffer.read(cx)) {
2286                    break;
2287                }
2288                next_message_ix += 1;
2289            }
2290
2291            let start = self.buffer.update(cx, |buffer, cx| {
2292                let offset = self
2293                    .message_anchors
2294                    .get(next_message_ix)
2295                    .map_or(buffer.len(), |message| {
2296                        buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
2297                    });
2298                buffer.edit([(offset..offset, "\n")], None, cx);
2299                buffer.anchor_before(offset + 1)
2300            });
2301
2302            let version = self.version.clone();
2303            let anchor = MessageAnchor {
2304                id: MessageId(self.next_timestamp()),
2305                start,
2306            };
2307            let metadata = MessageMetadata {
2308                role,
2309                status,
2310                timestamp: anchor.id.0,
2311                cache: None,
2312            };
2313            self.insert_message(anchor.clone(), metadata.clone(), cx);
2314            self.push_op(
2315                ContextOperation::InsertMessage {
2316                    anchor: anchor.clone(),
2317                    metadata,
2318                    version,
2319                },
2320                cx,
2321            );
2322            Some(anchor)
2323        } else {
2324            None
2325        }
2326    }
2327
2328    pub fn insert_content(&mut self, content: Content, cx: &mut ModelContext<Self>) {
2329        let buffer = self.buffer.read(cx);
2330        let insertion_ix = match self
2331            .contents
2332            .binary_search_by(|probe| probe.cmp(&content, buffer))
2333        {
2334            Ok(ix) => {
2335                self.contents.remove(ix);
2336                ix
2337            }
2338            Err(ix) => ix,
2339        };
2340        self.contents.insert(insertion_ix, content);
2341        cx.emit(ContextEvent::MessagesEdited);
2342    }
2343
2344    pub fn contents<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Content> {
2345        let buffer = self.buffer.read(cx);
2346        self.contents
2347            .iter()
2348            .filter(|content| {
2349                let range = content.range();
2350                range.start.is_valid(buffer) && range.end.is_valid(buffer)
2351            })
2352            .cloned()
2353    }
2354
2355    pub fn split_message(
2356        &mut self,
2357        range: Range<usize>,
2358        cx: &mut ModelContext<Self>,
2359    ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
2360        let start_message = self.message_for_offset(range.start, cx);
2361        let end_message = self.message_for_offset(range.end, cx);
2362        if let Some((start_message, end_message)) = start_message.zip(end_message) {
2363            // Prevent splitting when range spans multiple messages.
2364            if start_message.id != end_message.id {
2365                return (None, None);
2366            }
2367
2368            let message = start_message;
2369            let role = message.role;
2370            let mut edited_buffer = false;
2371
2372            let mut suffix_start = None;
2373
2374            // TODO: why did this start panicking?
2375            if range.start > message.offset_range.start
2376                && range.end < message.offset_range.end.saturating_sub(1)
2377            {
2378                if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
2379                    suffix_start = Some(range.end + 1);
2380                } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
2381                    suffix_start = Some(range.end);
2382                }
2383            }
2384
2385            let version = self.version.clone();
2386            let suffix = if let Some(suffix_start) = suffix_start {
2387                MessageAnchor {
2388                    id: MessageId(self.next_timestamp()),
2389                    start: self.buffer.read(cx).anchor_before(suffix_start),
2390                }
2391            } else {
2392                self.buffer.update(cx, |buffer, cx| {
2393                    buffer.edit([(range.end..range.end, "\n")], None, cx);
2394                });
2395                edited_buffer = true;
2396                MessageAnchor {
2397                    id: MessageId(self.next_timestamp()),
2398                    start: self.buffer.read(cx).anchor_before(range.end + 1),
2399                }
2400            };
2401
2402            let suffix_metadata = MessageMetadata {
2403                role,
2404                status: MessageStatus::Done,
2405                timestamp: suffix.id.0,
2406                cache: None,
2407            };
2408            self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
2409            self.push_op(
2410                ContextOperation::InsertMessage {
2411                    anchor: suffix.clone(),
2412                    metadata: suffix_metadata,
2413                    version,
2414                },
2415                cx,
2416            );
2417
2418            let new_messages =
2419                if range.start == range.end || range.start == message.offset_range.start {
2420                    (None, Some(suffix))
2421                } else {
2422                    let mut prefix_end = None;
2423                    if range.start > message.offset_range.start
2424                        && range.end < message.offset_range.end - 1
2425                    {
2426                        if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
2427                            prefix_end = Some(range.start + 1);
2428                        } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
2429                            == Some('\n')
2430                        {
2431                            prefix_end = Some(range.start);
2432                        }
2433                    }
2434
2435                    let version = self.version.clone();
2436                    let selection = if let Some(prefix_end) = prefix_end {
2437                        MessageAnchor {
2438                            id: MessageId(self.next_timestamp()),
2439                            start: self.buffer.read(cx).anchor_before(prefix_end),
2440                        }
2441                    } else {
2442                        self.buffer.update(cx, |buffer, cx| {
2443                            buffer.edit([(range.start..range.start, "\n")], None, cx)
2444                        });
2445                        edited_buffer = true;
2446                        MessageAnchor {
2447                            id: MessageId(self.next_timestamp()),
2448                            start: self.buffer.read(cx).anchor_before(range.end + 1),
2449                        }
2450                    };
2451
2452                    let selection_metadata = MessageMetadata {
2453                        role,
2454                        status: MessageStatus::Done,
2455                        timestamp: selection.id.0,
2456                        cache: None,
2457                    };
2458                    self.insert_message(selection.clone(), selection_metadata.clone(), cx);
2459                    self.push_op(
2460                        ContextOperation::InsertMessage {
2461                            anchor: selection.clone(),
2462                            metadata: selection_metadata,
2463                            version,
2464                        },
2465                        cx,
2466                    );
2467
2468                    (Some(selection), Some(suffix))
2469                };
2470
2471            if !edited_buffer {
2472                cx.emit(ContextEvent::MessagesEdited);
2473            }
2474            new_messages
2475        } else {
2476            (None, None)
2477        }
2478    }
2479
2480    fn insert_message(
2481        &mut self,
2482        new_anchor: MessageAnchor,
2483        new_metadata: MessageMetadata,
2484        cx: &mut ModelContext<Self>,
2485    ) {
2486        cx.emit(ContextEvent::MessagesEdited);
2487
2488        self.messages_metadata.insert(new_anchor.id, new_metadata);
2489
2490        let buffer = self.buffer.read(cx);
2491        let insertion_ix = self
2492            .message_anchors
2493            .iter()
2494            .position(|anchor| {
2495                let comparison = new_anchor.start.cmp(&anchor.start, buffer);
2496                comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
2497            })
2498            .unwrap_or(self.message_anchors.len());
2499        self.message_anchors.insert(insertion_ix, new_anchor);
2500    }
2501
2502    pub(super) fn summarize(&mut self, replace_old: bool, cx: &mut ModelContext<Self>) {
2503        let Some(provider) = LanguageModelRegistry::read_global(cx).active_provider() else {
2504            return;
2505        };
2506        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
2507            return;
2508        };
2509
2510        if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_none()) {
2511            if !provider.is_authenticated(cx) {
2512                return;
2513            }
2514
2515            let mut request = self.to_completion_request(RequestType::Chat, cx);
2516            request.messages.push(LanguageModelRequestMessage {
2517                role: Role::User,
2518                content: vec![
2519                    "Generate a concise 3-7 word title for this conversation, omitting punctuation"
2520                        .into(),
2521                ],
2522                cache: false,
2523            });
2524
2525            self.pending_summary = cx.spawn(|this, mut cx| {
2526                async move {
2527                    let stream = model.stream_completion_text(request, &cx);
2528                    let mut messages = stream.await?;
2529
2530                    let mut replaced = !replace_old;
2531                    while let Some(message) = messages.next().await {
2532                        let text = message?;
2533                        let mut lines = text.lines();
2534                        this.update(&mut cx, |this, cx| {
2535                            let version = this.version.clone();
2536                            let timestamp = this.next_timestamp();
2537                            let summary = this.summary.get_or_insert(ContextSummary::default());
2538                            if !replaced && replace_old {
2539                                summary.text.clear();
2540                                replaced = true;
2541                            }
2542                            summary.text.extend(lines.next());
2543                            summary.timestamp = timestamp;
2544                            let operation = ContextOperation::UpdateSummary {
2545                                summary: summary.clone(),
2546                                version,
2547                            };
2548                            this.push_op(operation, cx);
2549                            cx.emit(ContextEvent::SummaryChanged);
2550                        })?;
2551
2552                        // Stop if the LLM generated multiple lines.
2553                        if lines.next().is_some() {
2554                            break;
2555                        }
2556                    }
2557
2558                    this.update(&mut cx, |this, cx| {
2559                        let version = this.version.clone();
2560                        let timestamp = this.next_timestamp();
2561                        if let Some(summary) = this.summary.as_mut() {
2562                            summary.done = true;
2563                            summary.timestamp = timestamp;
2564                            let operation = ContextOperation::UpdateSummary {
2565                                summary: summary.clone(),
2566                                version,
2567                            };
2568                            this.push_op(operation, cx);
2569                            cx.emit(ContextEvent::SummaryChanged);
2570                        }
2571                    })?;
2572
2573                    anyhow::Ok(())
2574                }
2575                .log_err()
2576            });
2577        }
2578    }
2579
2580    fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
2581        self.messages_for_offsets([offset], cx).pop()
2582    }
2583
2584    pub fn messages_for_offsets(
2585        &self,
2586        offsets: impl IntoIterator<Item = usize>,
2587        cx: &AppContext,
2588    ) -> Vec<Message> {
2589        let mut result = Vec::new();
2590
2591        let mut messages = self.messages(cx).peekable();
2592        let mut offsets = offsets.into_iter().peekable();
2593        let mut current_message = messages.next();
2594        while let Some(offset) = offsets.next() {
2595            // Locate the message that contains the offset.
2596            while current_message.as_ref().map_or(false, |message| {
2597                !message.offset_range.contains(&offset) && messages.peek().is_some()
2598            }) {
2599                current_message = messages.next();
2600            }
2601            let Some(message) = current_message.as_ref() else {
2602                break;
2603            };
2604
2605            // Skip offsets that are in the same message.
2606            while offsets.peek().map_or(false, |offset| {
2607                message.offset_range.contains(offset) || messages.peek().is_none()
2608            }) {
2609                offsets.next();
2610            }
2611
2612            result.push(message.clone());
2613        }
2614        result
2615    }
2616
2617    fn messages_from_anchors<'a>(
2618        &'a self,
2619        message_anchors: impl Iterator<Item = &'a MessageAnchor> + 'a,
2620        cx: &'a AppContext,
2621    ) -> impl 'a + Iterator<Item = Message> {
2622        let buffer = self.buffer.read(cx);
2623
2624        Self::messages_from_iters(buffer, &self.messages_metadata, message_anchors.enumerate())
2625    }
2626
2627    pub fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
2628        self.messages_from_anchors(self.message_anchors.iter(), cx)
2629    }
2630
2631    pub fn messages_from_iters<'a>(
2632        buffer: &'a Buffer,
2633        metadata: &'a HashMap<MessageId, MessageMetadata>,
2634        messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
2635    ) -> impl 'a + Iterator<Item = Message> {
2636        let mut messages = messages.peekable();
2637
2638        iter::from_fn(move || {
2639            if let Some((start_ix, message_anchor)) = messages.next() {
2640                let metadata = metadata.get(&message_anchor.id)?;
2641
2642                let message_start = message_anchor.start.to_offset(buffer);
2643                let mut message_end = None;
2644                let mut end_ix = start_ix;
2645                while let Some((_, next_message)) = messages.peek() {
2646                    if next_message.start.is_valid(buffer) {
2647                        message_end = Some(next_message.start);
2648                        break;
2649                    } else {
2650                        end_ix += 1;
2651                        messages.next();
2652                    }
2653                }
2654                let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
2655                let message_end = message_end_anchor.to_offset(buffer);
2656
2657                return Some(Message {
2658                    index_range: start_ix..end_ix,
2659                    offset_range: message_start..message_end,
2660                    anchor_range: message_anchor.start..message_end_anchor,
2661                    id: message_anchor.id,
2662                    role: metadata.role,
2663                    status: metadata.status.clone(),
2664                    cache: metadata.cache.clone(),
2665                });
2666            }
2667            None
2668        })
2669    }
2670
2671    pub fn save(
2672        &mut self,
2673        debounce: Option<Duration>,
2674        fs: Arc<dyn Fs>,
2675        cx: &mut ModelContext<Context>,
2676    ) {
2677        if self.replica_id() != ReplicaId::default() {
2678            // Prevent saving a remote context for now.
2679            return;
2680        }
2681
2682        self.pending_save = cx.spawn(|this, mut cx| async move {
2683            if let Some(debounce) = debounce {
2684                cx.background_executor().timer(debounce).await;
2685            }
2686
2687            let (old_path, summary) = this.read_with(&cx, |this, _| {
2688                let path = this.path.clone();
2689                let summary = if let Some(summary) = this.summary.as_ref() {
2690                    if summary.done {
2691                        Some(summary.text.clone())
2692                    } else {
2693                        None
2694                    }
2695                } else {
2696                    None
2697                };
2698                (path, summary)
2699            })?;
2700
2701            if let Some(summary) = summary {
2702                let context = this.read_with(&cx, |this, cx| this.serialize(cx))?;
2703                let mut discriminant = 1;
2704                let mut new_path;
2705                loop {
2706                    new_path = contexts_dir().join(&format!(
2707                        "{} - {}.zed.json",
2708                        summary.trim(),
2709                        discriminant
2710                    ));
2711                    if fs.is_file(&new_path).await {
2712                        discriminant += 1;
2713                    } else {
2714                        break;
2715                    }
2716                }
2717
2718                fs.create_dir(contexts_dir().as_ref()).await?;
2719                fs.atomic_write(new_path.clone(), serde_json::to_string(&context).unwrap())
2720                    .await?;
2721                if let Some(old_path) = old_path {
2722                    if new_path != old_path {
2723                        fs.remove_file(
2724                            &old_path,
2725                            RemoveOptions {
2726                                recursive: false,
2727                                ignore_if_not_exists: true,
2728                            },
2729                        )
2730                        .await?;
2731                    }
2732                }
2733
2734                this.update(&mut cx, |this, _| this.path = Some(new_path))?;
2735            }
2736
2737            Ok(())
2738        });
2739    }
2740
2741    pub(crate) fn custom_summary(&mut self, custom_summary: String, cx: &mut ModelContext<Self>) {
2742        let timestamp = self.next_timestamp();
2743        let summary = self.summary.get_or_insert(ContextSummary::default());
2744        summary.timestamp = timestamp;
2745        summary.done = true;
2746        summary.text = custom_summary;
2747        cx.emit(ContextEvent::SummaryChanged);
2748    }
2749}
2750
2751fn trimmed_text_in_range(buffer: &BufferSnapshot, range: Range<text::Anchor>) -> String {
2752    let mut is_start = true;
2753    let mut content = buffer
2754        .text_for_range(range)
2755        .map(|mut chunk| {
2756            if is_start {
2757                chunk = chunk.trim_start_matches('\n');
2758                if !chunk.is_empty() {
2759                    is_start = false;
2760                }
2761            }
2762            chunk
2763        })
2764        .collect::<String>();
2765    content.truncate(content.trim_end().len());
2766    content
2767}
2768
2769#[derive(Debug, Default)]
2770pub struct ContextVersion {
2771    context: clock::Global,
2772    buffer: clock::Global,
2773}
2774
2775impl ContextVersion {
2776    pub fn from_proto(proto: &proto::ContextVersion) -> Self {
2777        Self {
2778            context: language::proto::deserialize_version(&proto.context_version),
2779            buffer: language::proto::deserialize_version(&proto.buffer_version),
2780        }
2781    }
2782
2783    pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
2784        proto::ContextVersion {
2785            context_id: context_id.to_proto(),
2786            context_version: language::proto::serialize_version(&self.context),
2787            buffer_version: language::proto::serialize_version(&self.buffer),
2788        }
2789    }
2790}
2791
2792#[derive(Debug, Clone)]
2793pub struct PendingSlashCommand {
2794    pub name: String,
2795    pub arguments: SmallVec<[String; 3]>,
2796    pub status: PendingSlashCommandStatus,
2797    pub source_range: Range<language::Anchor>,
2798}
2799
2800#[derive(Debug, Clone)]
2801pub enum PendingSlashCommandStatus {
2802    Idle,
2803    Running { _task: Shared<Task<()>> },
2804    Error(String),
2805}
2806
2807pub(crate) struct ToolUseFeatureFlag;
2808
2809impl FeatureFlag for ToolUseFeatureFlag {
2810    const NAME: &'static str = "assistant-tool-use";
2811
2812    fn enabled_for_staff() -> bool {
2813        false
2814    }
2815}
2816
2817#[derive(Debug, Clone)]
2818pub struct PendingToolUse {
2819    pub id: Arc<str>,
2820    pub name: String,
2821    pub input: serde_json::Value,
2822    pub status: PendingToolUseStatus,
2823    pub source_range: Range<language::Anchor>,
2824}
2825
2826#[derive(Debug, Clone)]
2827pub enum PendingToolUseStatus {
2828    Idle,
2829    Running { _task: Shared<Task<()>> },
2830    Error(String),
2831}
2832
2833impl PendingToolUseStatus {
2834    pub fn is_idle(&self) -> bool {
2835        matches!(self, PendingToolUseStatus::Idle)
2836    }
2837}
2838
2839#[derive(Serialize, Deserialize)]
2840pub struct SavedMessage {
2841    pub id: MessageId,
2842    pub start: usize,
2843    pub metadata: MessageMetadata,
2844}
2845
2846#[derive(Serialize, Deserialize)]
2847pub struct SavedContext {
2848    pub id: Option<ContextId>,
2849    pub zed: String,
2850    pub version: String,
2851    pub text: String,
2852    pub messages: Vec<SavedMessage>,
2853    pub summary: String,
2854    pub slash_command_output_sections:
2855        Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
2856}
2857
2858impl SavedContext {
2859    pub const VERSION: &'static str = "0.4.0";
2860
2861    pub fn from_json(json: &str) -> Result<Self> {
2862        let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
2863        match saved_context_json
2864            .get("version")
2865            .ok_or_else(|| anyhow!("version not found"))?
2866        {
2867            serde_json::Value::String(version) => match version.as_str() {
2868                SavedContext::VERSION => {
2869                    Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
2870                }
2871                SavedContextV0_3_0::VERSION => {
2872                    let saved_context =
2873                        serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
2874                    Ok(saved_context.upgrade())
2875                }
2876                SavedContextV0_2_0::VERSION => {
2877                    let saved_context =
2878                        serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
2879                    Ok(saved_context.upgrade())
2880                }
2881                SavedContextV0_1_0::VERSION => {
2882                    let saved_context =
2883                        serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
2884                    Ok(saved_context.upgrade())
2885                }
2886                _ => Err(anyhow!("unrecognized saved context version: {}", version)),
2887            },
2888            _ => Err(anyhow!("version not found on saved context")),
2889        }
2890    }
2891
2892    fn into_ops(
2893        self,
2894        buffer: &Model<Buffer>,
2895        cx: &mut ModelContext<Context>,
2896    ) -> Vec<ContextOperation> {
2897        let mut operations = Vec::new();
2898        let mut version = clock::Global::new();
2899        let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
2900
2901        let mut first_message_metadata = None;
2902        for message in self.messages {
2903            if message.id == MessageId(clock::Lamport::default()) {
2904                first_message_metadata = Some(message.metadata);
2905            } else {
2906                operations.push(ContextOperation::InsertMessage {
2907                    anchor: MessageAnchor {
2908                        id: message.id,
2909                        start: buffer.read(cx).anchor_before(message.start),
2910                    },
2911                    metadata: MessageMetadata {
2912                        role: message.metadata.role,
2913                        status: message.metadata.status,
2914                        timestamp: message.metadata.timestamp,
2915                        cache: None,
2916                    },
2917                    version: version.clone(),
2918                });
2919                version.observe(message.id.0);
2920                next_timestamp.observe(message.id.0);
2921            }
2922        }
2923
2924        if let Some(metadata) = first_message_metadata {
2925            let timestamp = next_timestamp.tick();
2926            operations.push(ContextOperation::UpdateMessage {
2927                message_id: MessageId(clock::Lamport::default()),
2928                metadata: MessageMetadata {
2929                    role: metadata.role,
2930                    status: metadata.status,
2931                    timestamp,
2932                    cache: None,
2933                },
2934                version: version.clone(),
2935            });
2936            version.observe(timestamp);
2937        }
2938
2939        let timestamp = next_timestamp.tick();
2940        operations.push(ContextOperation::SlashCommandFinished {
2941            id: SlashCommandId(timestamp),
2942            output_range: language::Anchor::MIN..language::Anchor::MAX,
2943            sections: self
2944                .slash_command_output_sections
2945                .into_iter()
2946                .map(|section| {
2947                    let buffer = buffer.read(cx);
2948                    SlashCommandOutputSection {
2949                        range: buffer.anchor_after(section.range.start)
2950                            ..buffer.anchor_before(section.range.end),
2951                        icon: section.icon,
2952                        label: section.label,
2953                        metadata: section.metadata,
2954                    }
2955                })
2956                .collect(),
2957            version: version.clone(),
2958        });
2959        version.observe(timestamp);
2960
2961        let timestamp = next_timestamp.tick();
2962        operations.push(ContextOperation::UpdateSummary {
2963            summary: ContextSummary {
2964                text: self.summary,
2965                done: true,
2966                timestamp,
2967            },
2968            version: version.clone(),
2969        });
2970        version.observe(timestamp);
2971
2972        operations
2973    }
2974}
2975
2976#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
2977struct SavedMessageIdPreV0_4_0(usize);
2978
2979#[derive(Serialize, Deserialize)]
2980struct SavedMessagePreV0_4_0 {
2981    id: SavedMessageIdPreV0_4_0,
2982    start: usize,
2983}
2984
2985#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
2986struct SavedMessageMetadataPreV0_4_0 {
2987    role: Role,
2988    status: MessageStatus,
2989}
2990
2991#[derive(Serialize, Deserialize)]
2992struct SavedContextV0_3_0 {
2993    id: Option<ContextId>,
2994    zed: String,
2995    version: String,
2996    text: String,
2997    messages: Vec<SavedMessagePreV0_4_0>,
2998    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
2999    summary: String,
3000    slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3001}
3002
3003impl SavedContextV0_3_0 {
3004    const VERSION: &'static str = "0.3.0";
3005
3006    fn upgrade(self) -> SavedContext {
3007        SavedContext {
3008            id: self.id,
3009            zed: self.zed,
3010            version: SavedContext::VERSION.into(),
3011            text: self.text,
3012            messages: self
3013                .messages
3014                .into_iter()
3015                .filter_map(|message| {
3016                    let metadata = self.message_metadata.get(&message.id)?;
3017                    let timestamp = clock::Lamport {
3018                        replica_id: ReplicaId::default(),
3019                        value: message.id.0 as u32,
3020                    };
3021                    Some(SavedMessage {
3022                        id: MessageId(timestamp),
3023                        start: message.start,
3024                        metadata: MessageMetadata {
3025                            role: metadata.role,
3026                            status: metadata.status.clone(),
3027                            timestamp,
3028                            cache: None,
3029                        },
3030                    })
3031                })
3032                .collect(),
3033            summary: self.summary,
3034            slash_command_output_sections: self.slash_command_output_sections,
3035        }
3036    }
3037}
3038
3039#[derive(Serialize, Deserialize)]
3040struct SavedContextV0_2_0 {
3041    id: Option<ContextId>,
3042    zed: String,
3043    version: String,
3044    text: String,
3045    messages: Vec<SavedMessagePreV0_4_0>,
3046    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3047    summary: String,
3048}
3049
3050impl SavedContextV0_2_0 {
3051    const VERSION: &'static str = "0.2.0";
3052
3053    fn upgrade(self) -> SavedContext {
3054        SavedContextV0_3_0 {
3055            id: self.id,
3056            zed: self.zed,
3057            version: SavedContextV0_3_0::VERSION.to_string(),
3058            text: self.text,
3059            messages: self.messages,
3060            message_metadata: self.message_metadata,
3061            summary: self.summary,
3062            slash_command_output_sections: Vec::new(),
3063        }
3064        .upgrade()
3065    }
3066}
3067
3068#[derive(Serialize, Deserialize)]
3069struct SavedContextV0_1_0 {
3070    id: Option<ContextId>,
3071    zed: String,
3072    version: String,
3073    text: String,
3074    messages: Vec<SavedMessagePreV0_4_0>,
3075    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3076    summary: String,
3077    api_url: Option<String>,
3078    model: OpenAiModel,
3079}
3080
3081impl SavedContextV0_1_0 {
3082    const VERSION: &'static str = "0.1.0";
3083
3084    fn upgrade(self) -> SavedContext {
3085        SavedContextV0_2_0 {
3086            id: self.id,
3087            zed: self.zed,
3088            version: SavedContextV0_2_0::VERSION.to_string(),
3089            text: self.text,
3090            messages: self.messages,
3091            message_metadata: self.message_metadata,
3092            summary: self.summary,
3093        }
3094        .upgrade()
3095    }
3096}
3097
3098#[derive(Clone)]
3099pub struct SavedContextMetadata {
3100    pub title: String,
3101    pub path: PathBuf,
3102    pub mtime: chrono::DateTime<chrono::Local>,
3103}