context.rs

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