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