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