context.rs

   1#[cfg(test)]
   2mod context_tests;
   3
   4use anyhow::{Context as _, Result, bail};
   5use assistant_settings::AssistantSettings;
   6use assistant_slash_command::{
   7    SlashCommandContent, SlashCommandEvent, SlashCommandLine, SlashCommandOutputSection,
   8    SlashCommandResult, SlashCommandWorkingSet,
   9};
  10use assistant_slash_commands::FileCommandMetadata;
  11use client::{self, proto, telemetry::Telemetry};
  12use clock::ReplicaId;
  13use collections::{HashMap, HashSet};
  14use fs::{Fs, RemoveOptions};
  15use futures::{FutureExt, StreamExt, future::Shared};
  16use gpui::{
  17    App, AppContext as _, Context, Entity, EventEmitter, RenderImage, SharedString, Subscription,
  18    Task,
  19};
  20use language::{AnchorRangeExt, Bias, Buffer, LanguageRegistry, OffsetRangeExt, Point, ToOffset};
  21use language_model::{
  22    LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionEvent,
  23    LanguageModelImage, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
  24    LanguageModelToolUseId, MessageContent, PaymentRequiredError, Role, StopReason,
  25    report_assistant_event,
  26};
  27use open_ai::Model as OpenAiModel;
  28use paths::contexts_dir;
  29use project::Project;
  30use prompt_store::PromptBuilder;
  31use serde::{Deserialize, Serialize};
  32use smallvec::SmallVec;
  33use std::{
  34    cmp::{Ordering, max},
  35    fmt::{Debug, Write as _},
  36    iter, mem,
  37    ops::Range,
  38    path::Path,
  39    sync::Arc,
  40    time::{Duration, Instant},
  41};
  42use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
  43use text::{BufferSnapshot, ToPoint};
  44use ui::IconName;
  45use util::{ResultExt, TryFutureExt, post_inc};
  46use uuid::Uuid;
  47use zed_llm_client::CompletionIntent;
  48
  49#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
  50pub struct ContextId(String);
  51
  52impl ContextId {
  53    pub fn new() -> Self {
  54        Self(Uuid::new_v4().to_string())
  55    }
  56
  57    pub fn from_proto(id: String) -> Self {
  58        Self(id)
  59    }
  60
  61    pub fn to_proto(&self) -> String {
  62        self.0.clone()
  63    }
  64}
  65
  66#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
  67pub struct MessageId(pub clock::Lamport);
  68
  69impl MessageId {
  70    pub fn as_u64(self) -> u64 {
  71        self.0.as_u64()
  72    }
  73}
  74
  75#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
  76pub enum MessageStatus {
  77    Pending,
  78    Done,
  79    Error(SharedString),
  80    Canceled,
  81}
  82
  83impl MessageStatus {
  84    pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
  85        match status.variant {
  86            Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
  87            Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
  88            Some(proto::context_message_status::Variant::Error(error)) => {
  89                MessageStatus::Error(error.message.into())
  90            }
  91            Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
  92            None => MessageStatus::Pending,
  93        }
  94    }
  95
  96    pub fn to_proto(&self) -> proto::ContextMessageStatus {
  97        match self {
  98            MessageStatus::Pending => proto::ContextMessageStatus {
  99                variant: Some(proto::context_message_status::Variant::Pending(
 100                    proto::context_message_status::Pending {},
 101                )),
 102            },
 103            MessageStatus::Done => proto::ContextMessageStatus {
 104                variant: Some(proto::context_message_status::Variant::Done(
 105                    proto::context_message_status::Done {},
 106                )),
 107            },
 108            MessageStatus::Error(message) => proto::ContextMessageStatus {
 109                variant: Some(proto::context_message_status::Variant::Error(
 110                    proto::context_message_status::Error {
 111                        message: message.to_string(),
 112                    },
 113                )),
 114            },
 115            MessageStatus::Canceled => proto::ContextMessageStatus {
 116                variant: Some(proto::context_message_status::Variant::Canceled(
 117                    proto::context_message_status::Canceled {},
 118                )),
 119            },
 120        }
 121    }
 122}
 123
 124#[derive(Clone, Debug)]
 125pub enum ContextOperation {
 126    InsertMessage {
 127        anchor: MessageAnchor,
 128        metadata: MessageMetadata,
 129        version: clock::Global,
 130    },
 131    UpdateMessage {
 132        message_id: MessageId,
 133        metadata: MessageMetadata,
 134        version: clock::Global,
 135    },
 136    UpdateSummary {
 137        summary: ContextSummaryContent,
 138        version: clock::Global,
 139    },
 140    SlashCommandStarted {
 141        id: InvokedSlashCommandId,
 142        output_range: Range<language::Anchor>,
 143        name: String,
 144        version: clock::Global,
 145    },
 146    SlashCommandFinished {
 147        id: InvokedSlashCommandId,
 148        timestamp: clock::Lamport,
 149        error_message: Option<String>,
 150        version: clock::Global,
 151    },
 152    SlashCommandOutputSectionAdded {
 153        timestamp: clock::Lamport,
 154        section: SlashCommandOutputSection<language::Anchor>,
 155        version: clock::Global,
 156    },
 157    ThoughtProcessOutputSectionAdded {
 158        timestamp: clock::Lamport,
 159        section: ThoughtProcessOutputSection<language::Anchor>,
 160        version: clock::Global,
 161    },
 162    BufferOperation(language::Operation),
 163}
 164
 165impl ContextOperation {
 166    pub fn from_proto(op: proto::ContextOperation) -> Result<Self> {
 167        match op.variant.context("invalid variant")? {
 168            proto::context_operation::Variant::InsertMessage(insert) => {
 169                let message = insert.message.context("invalid message")?;
 170                let id = MessageId(language::proto::deserialize_timestamp(
 171                    message.id.context("invalid id")?,
 172                ));
 173                Ok(Self::InsertMessage {
 174                    anchor: MessageAnchor {
 175                        id,
 176                        start: language::proto::deserialize_anchor(
 177                            message.start.context("invalid anchor")?,
 178                        )
 179                        .context("invalid anchor")?,
 180                    },
 181                    metadata: MessageMetadata {
 182                        role: Role::from_proto(message.role),
 183                        status: MessageStatus::from_proto(
 184                            message.status.context("invalid status")?,
 185                        ),
 186                        timestamp: id.0,
 187                        cache: None,
 188                    },
 189                    version: language::proto::deserialize_version(&insert.version),
 190                })
 191            }
 192            proto::context_operation::Variant::UpdateMessage(update) => Ok(Self::UpdateMessage {
 193                message_id: MessageId(language::proto::deserialize_timestamp(
 194                    update.message_id.context("invalid message id")?,
 195                )),
 196                metadata: MessageMetadata {
 197                    role: Role::from_proto(update.role),
 198                    status: MessageStatus::from_proto(update.status.context("invalid status")?),
 199                    timestamp: language::proto::deserialize_timestamp(
 200                        update.timestamp.context("invalid timestamp")?,
 201                    ),
 202                    cache: None,
 203                },
 204                version: language::proto::deserialize_version(&update.version),
 205            }),
 206            proto::context_operation::Variant::UpdateSummary(update) => Ok(Self::UpdateSummary {
 207                summary: ContextSummaryContent {
 208                    text: update.summary,
 209                    done: update.done,
 210                    timestamp: language::proto::deserialize_timestamp(
 211                        update.timestamp.context("invalid timestamp")?,
 212                    ),
 213                },
 214                version: language::proto::deserialize_version(&update.version),
 215            }),
 216            proto::context_operation::Variant::SlashCommandStarted(message) => {
 217                Ok(Self::SlashCommandStarted {
 218                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
 219                        message.id.context("invalid id")?,
 220                    )),
 221                    output_range: language::proto::deserialize_anchor_range(
 222                        message.output_range.context("invalid range")?,
 223                    )?,
 224                    name: message.name,
 225                    version: language::proto::deserialize_version(&message.version),
 226                })
 227            }
 228            proto::context_operation::Variant::SlashCommandOutputSectionAdded(message) => {
 229                let section = message.section.context("missing section")?;
 230                Ok(Self::SlashCommandOutputSectionAdded {
 231                    timestamp: language::proto::deserialize_timestamp(
 232                        message.timestamp.context("missing timestamp")?,
 233                    ),
 234                    section: SlashCommandOutputSection {
 235                        range: language::proto::deserialize_anchor_range(
 236                            section.range.context("invalid range")?,
 237                        )?,
 238                        icon: section.icon_name.parse()?,
 239                        label: section.label.into(),
 240                        metadata: section
 241                            .metadata
 242                            .and_then(|metadata| serde_json::from_str(&metadata).log_err()),
 243                    },
 244                    version: language::proto::deserialize_version(&message.version),
 245                })
 246            }
 247            proto::context_operation::Variant::SlashCommandCompleted(message) => {
 248                Ok(Self::SlashCommandFinished {
 249                    id: InvokedSlashCommandId(language::proto::deserialize_timestamp(
 250                        message.id.context("invalid id")?,
 251                    )),
 252                    timestamp: language::proto::deserialize_timestamp(
 253                        message.timestamp.context("missing timestamp")?,
 254                    ),
 255                    error_message: message.error_message,
 256                    version: language::proto::deserialize_version(&message.version),
 257                })
 258            }
 259            proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(message) => {
 260                let section = message.section.context("missing section")?;
 261                Ok(Self::ThoughtProcessOutputSectionAdded {
 262                    timestamp: language::proto::deserialize_timestamp(
 263                        message.timestamp.context("missing timestamp")?,
 264                    ),
 265                    section: ThoughtProcessOutputSection {
 266                        range: language::proto::deserialize_anchor_range(
 267                            section.range.context("invalid range")?,
 268                        )?,
 269                    },
 270                    version: language::proto::deserialize_version(&message.version),
 271                })
 272            }
 273            proto::context_operation::Variant::BufferOperation(op) => Ok(Self::BufferOperation(
 274                language::proto::deserialize_operation(
 275                    op.operation.context("invalid buffer operation")?,
 276                )?,
 277            )),
 278        }
 279    }
 280
 281    pub fn to_proto(&self) -> proto::ContextOperation {
 282        match self {
 283            Self::InsertMessage {
 284                anchor,
 285                metadata,
 286                version,
 287            } => proto::ContextOperation {
 288                variant: Some(proto::context_operation::Variant::InsertMessage(
 289                    proto::context_operation::InsertMessage {
 290                        message: Some(proto::ContextMessage {
 291                            id: Some(language::proto::serialize_timestamp(anchor.id.0)),
 292                            start: Some(language::proto::serialize_anchor(&anchor.start)),
 293                            role: metadata.role.to_proto() as i32,
 294                            status: Some(metadata.status.to_proto()),
 295                        }),
 296                        version: language::proto::serialize_version(version),
 297                    },
 298                )),
 299            },
 300            Self::UpdateMessage {
 301                message_id,
 302                metadata,
 303                version,
 304            } => proto::ContextOperation {
 305                variant: Some(proto::context_operation::Variant::UpdateMessage(
 306                    proto::context_operation::UpdateMessage {
 307                        message_id: Some(language::proto::serialize_timestamp(message_id.0)),
 308                        role: metadata.role.to_proto() as i32,
 309                        status: Some(metadata.status.to_proto()),
 310                        timestamp: Some(language::proto::serialize_timestamp(metadata.timestamp)),
 311                        version: language::proto::serialize_version(version),
 312                    },
 313                )),
 314            },
 315            Self::UpdateSummary { summary, version } => proto::ContextOperation {
 316                variant: Some(proto::context_operation::Variant::UpdateSummary(
 317                    proto::context_operation::UpdateSummary {
 318                        summary: summary.text.clone(),
 319                        done: summary.done,
 320                        timestamp: Some(language::proto::serialize_timestamp(summary.timestamp)),
 321                        version: language::proto::serialize_version(version),
 322                    },
 323                )),
 324            },
 325            Self::SlashCommandStarted {
 326                id,
 327                output_range,
 328                name,
 329                version,
 330            } => proto::ContextOperation {
 331                variant: Some(proto::context_operation::Variant::SlashCommandStarted(
 332                    proto::context_operation::SlashCommandStarted {
 333                        id: Some(language::proto::serialize_timestamp(id.0)),
 334                        output_range: Some(language::proto::serialize_anchor_range(
 335                            output_range.clone(),
 336                        )),
 337                        name: name.clone(),
 338                        version: language::proto::serialize_version(version),
 339                    },
 340                )),
 341            },
 342            Self::SlashCommandOutputSectionAdded {
 343                timestamp,
 344                section,
 345                version,
 346            } => proto::ContextOperation {
 347                variant: Some(
 348                    proto::context_operation::Variant::SlashCommandOutputSectionAdded(
 349                        proto::context_operation::SlashCommandOutputSectionAdded {
 350                            timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 351                            section: Some({
 352                                let icon_name: &'static str = section.icon.into();
 353                                proto::SlashCommandOutputSection {
 354                                    range: Some(language::proto::serialize_anchor_range(
 355                                        section.range.clone(),
 356                                    )),
 357                                    icon_name: icon_name.to_string(),
 358                                    label: section.label.to_string(),
 359                                    metadata: section.metadata.as_ref().and_then(|metadata| {
 360                                        serde_json::to_string(metadata).log_err()
 361                                    }),
 362                                }
 363                            }),
 364                            version: language::proto::serialize_version(version),
 365                        },
 366                    ),
 367                ),
 368            },
 369            Self::SlashCommandFinished {
 370                id,
 371                timestamp,
 372                error_message,
 373                version,
 374            } => proto::ContextOperation {
 375                variant: Some(proto::context_operation::Variant::SlashCommandCompleted(
 376                    proto::context_operation::SlashCommandCompleted {
 377                        id: Some(language::proto::serialize_timestamp(id.0)),
 378                        timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 379                        error_message: error_message.clone(),
 380                        version: language::proto::serialize_version(version),
 381                    },
 382                )),
 383            },
 384            Self::ThoughtProcessOutputSectionAdded {
 385                timestamp,
 386                section,
 387                version,
 388            } => proto::ContextOperation {
 389                variant: Some(
 390                    proto::context_operation::Variant::ThoughtProcessOutputSectionAdded(
 391                        proto::context_operation::ThoughtProcessOutputSectionAdded {
 392                            timestamp: Some(language::proto::serialize_timestamp(*timestamp)),
 393                            section: Some({
 394                                proto::ThoughtProcessOutputSection {
 395                                    range: Some(language::proto::serialize_anchor_range(
 396                                        section.range.clone(),
 397                                    )),
 398                                }
 399                            }),
 400                            version: language::proto::serialize_version(version),
 401                        },
 402                    ),
 403                ),
 404            },
 405            Self::BufferOperation(operation) => proto::ContextOperation {
 406                variant: Some(proto::context_operation::Variant::BufferOperation(
 407                    proto::context_operation::BufferOperation {
 408                        operation: Some(language::proto::serialize_operation(operation)),
 409                    },
 410                )),
 411            },
 412        }
 413    }
 414
 415    fn timestamp(&self) -> clock::Lamport {
 416        match self {
 417            Self::InsertMessage { anchor, .. } => anchor.id.0,
 418            Self::UpdateMessage { metadata, .. } => metadata.timestamp,
 419            Self::UpdateSummary { summary, .. } => summary.timestamp,
 420            Self::SlashCommandStarted { id, .. } => id.0,
 421            Self::SlashCommandOutputSectionAdded { timestamp, .. }
 422            | Self::SlashCommandFinished { timestamp, .. }
 423            | Self::ThoughtProcessOutputSectionAdded { timestamp, .. } => *timestamp,
 424            Self::BufferOperation(_) => {
 425                panic!("reading the timestamp of a buffer operation is not supported")
 426            }
 427        }
 428    }
 429
 430    /// Returns the current version of the context operation.
 431    pub fn version(&self) -> &clock::Global {
 432        match self {
 433            Self::InsertMessage { version, .. }
 434            | Self::UpdateMessage { version, .. }
 435            | Self::UpdateSummary { version, .. }
 436            | Self::SlashCommandStarted { version, .. }
 437            | Self::SlashCommandOutputSectionAdded { version, .. }
 438            | Self::SlashCommandFinished { version, .. }
 439            | Self::ThoughtProcessOutputSectionAdded { version, .. } => version,
 440            Self::BufferOperation(_) => {
 441                panic!("reading the version of a buffer operation is not supported")
 442            }
 443        }
 444    }
 445}
 446
 447#[derive(Debug, Clone)]
 448pub enum ContextEvent {
 449    ShowAssistError(SharedString),
 450    ShowPaymentRequiredError,
 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 {
2159                            let error_message = error
2160                                .chain()
2161                                .map(|err| err.to_string())
2162                                .collect::<Vec<_>>()
2163                                .join("\n");
2164                            cx.emit(ContextEvent::ShowAssistError(SharedString::from(
2165                                error_message.clone(),
2166                            )));
2167                            this.update_metadata(assistant_message_id, cx, |metadata| {
2168                                metadata.status =
2169                                    MessageStatus::Error(SharedString::from(error_message.clone()));
2170                            });
2171                            Some(error_message)
2172                        }
2173                    } else {
2174                        this.update_metadata(assistant_message_id, cx, |metadata| {
2175                            metadata.status = MessageStatus::Done;
2176                        });
2177                        None
2178                    };
2179
2180                    let language_name = this
2181                        .buffer
2182                        .read(cx)
2183                        .language()
2184                        .map(|language| language.name());
2185                    report_assistant_event(
2186                        AssistantEventData {
2187                            conversation_id: Some(this.id.0.clone()),
2188                            kind: AssistantKind::Panel,
2189                            phase: AssistantPhase::Response,
2190                            message_id: None,
2191                            model: model.telemetry_id(),
2192                            model_provider: model.provider_id().to_string(),
2193                            response_latency,
2194                            error_message,
2195                            language_name: language_name.map(|name| name.to_proto()),
2196                        },
2197                        this.telemetry.clone(),
2198                        cx.http_client(),
2199                        model.api_key(cx),
2200                        cx.background_executor(),
2201                    );
2202
2203                    if let Ok(stop_reason) = result {
2204                        match stop_reason {
2205                            StopReason::ToolUse => {}
2206                            StopReason::EndTurn => {}
2207                            StopReason::MaxTokens => {}
2208                            StopReason::Refusal => {}
2209                        }
2210                    }
2211                })
2212                .ok();
2213            }
2214        });
2215
2216        self.pending_completions.push(PendingCompletion {
2217            id: pending_completion_id,
2218            assistant_message_id: assistant_message.id,
2219            _task: task,
2220        });
2221
2222        Some(user_message)
2223    }
2224
2225    pub fn to_xml(&self, cx: &App) -> String {
2226        let mut output = String::new();
2227        let buffer = self.buffer.read(cx);
2228        for message in self.messages(cx) {
2229            if message.status != MessageStatus::Done {
2230                continue;
2231            }
2232
2233            writeln!(&mut output, "<{}>", message.role).unwrap();
2234            for chunk in buffer.text_for_range(message.offset_range) {
2235                output.push_str(chunk);
2236            }
2237            if !output.ends_with('\n') {
2238                output.push('\n');
2239            }
2240            writeln!(&mut output, "</{}>", message.role).unwrap();
2241        }
2242        output
2243    }
2244
2245    pub fn to_completion_request(
2246        &self,
2247        model: Option<&Arc<dyn LanguageModel>>,
2248        cx: &App,
2249    ) -> LanguageModelRequest {
2250        let buffer = self.buffer.read(cx);
2251
2252        let mut contents = self.contents(cx).peekable();
2253
2254        fn collect_text_content(buffer: &Buffer, range: Range<usize>) -> Option<String> {
2255            let text: String = buffer.text_for_range(range.clone()).collect();
2256            if text.trim().is_empty() {
2257                None
2258            } else {
2259                Some(text)
2260            }
2261        }
2262
2263        let mut completion_request = LanguageModelRequest {
2264            thread_id: None,
2265            prompt_id: None,
2266            intent: Some(CompletionIntent::UserPrompt),
2267            mode: None,
2268            messages: Vec::new(),
2269            tools: Vec::new(),
2270            tool_choice: None,
2271            stop: Vec::new(),
2272            temperature: model
2273                .and_then(|model| AssistantSettings::temperature_for_model(model, cx)),
2274        };
2275        for message in self.messages(cx) {
2276            if message.status != MessageStatus::Done {
2277                continue;
2278            }
2279
2280            let mut offset = message.offset_range.start;
2281            let mut request_message = LanguageModelRequestMessage {
2282                role: message.role,
2283                content: Vec::new(),
2284                cache: message
2285                    .cache
2286                    .as_ref()
2287                    .map_or(false, |cache| cache.is_anchor),
2288            };
2289
2290            while let Some(content) = contents.peek() {
2291                if content
2292                    .range()
2293                    .end
2294                    .cmp(&message.anchor_range.end, buffer)
2295                    .is_lt()
2296                {
2297                    let content = contents.next().unwrap();
2298                    let range = content.range().to_offset(buffer);
2299                    request_message.content.extend(
2300                        collect_text_content(buffer, offset..range.start).map(MessageContent::Text),
2301                    );
2302
2303                    match content {
2304                        Content::Image { image, .. } => {
2305                            if let Some(image) = image.clone().now_or_never().flatten() {
2306                                request_message
2307                                    .content
2308                                    .push(language_model::MessageContent::Image(image));
2309                            }
2310                        }
2311                    }
2312
2313                    offset = range.end;
2314                } else {
2315                    break;
2316                }
2317            }
2318
2319            request_message.content.extend(
2320                collect_text_content(buffer, offset..message.offset_range.end)
2321                    .map(MessageContent::Text),
2322            );
2323
2324            if !request_message.contents_empty() {
2325                completion_request.messages.push(request_message);
2326            }
2327        }
2328
2329        completion_request
2330    }
2331
2332    pub fn cancel_last_assist(&mut self, cx: &mut Context<Self>) -> bool {
2333        if let Some(pending_completion) = self.pending_completions.pop() {
2334            self.update_metadata(pending_completion.assistant_message_id, cx, |metadata| {
2335                if metadata.status == MessageStatus::Pending {
2336                    metadata.status = MessageStatus::Canceled;
2337                }
2338            });
2339            true
2340        } else {
2341            false
2342        }
2343    }
2344
2345    pub fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2346        for id in &ids {
2347            if let Some(metadata) = self.messages_metadata.get(id) {
2348                let role = metadata.role.cycle();
2349                self.update_metadata(*id, cx, |metadata| metadata.role = role);
2350            }
2351        }
2352
2353        self.message_roles_updated(ids, cx);
2354    }
2355
2356    fn message_roles_updated(&mut self, ids: HashSet<MessageId>, cx: &mut Context<Self>) {
2357        let mut ranges = Vec::new();
2358        for message in self.messages(cx) {
2359            if ids.contains(&message.id) {
2360                ranges.push(message.anchor_range.clone());
2361            }
2362        }
2363    }
2364
2365    pub fn update_metadata(
2366        &mut self,
2367        id: MessageId,
2368        cx: &mut Context<Self>,
2369        f: impl FnOnce(&mut MessageMetadata),
2370    ) {
2371        let version = self.version.clone();
2372        let timestamp = self.next_timestamp();
2373        if let Some(metadata) = self.messages_metadata.get_mut(&id) {
2374            f(metadata);
2375            metadata.timestamp = timestamp;
2376            let operation = ContextOperation::UpdateMessage {
2377                message_id: id,
2378                metadata: metadata.clone(),
2379                version,
2380            };
2381            self.push_op(operation, cx);
2382            cx.emit(ContextEvent::MessagesEdited);
2383            cx.notify();
2384        }
2385    }
2386
2387    pub fn insert_message_after(
2388        &mut self,
2389        message_id: MessageId,
2390        role: Role,
2391        status: MessageStatus,
2392        cx: &mut Context<Self>,
2393    ) -> Option<MessageAnchor> {
2394        if let Some(prev_message_ix) = self
2395            .message_anchors
2396            .iter()
2397            .position(|message| message.id == message_id)
2398        {
2399            // Find the next valid message after the one we were given.
2400            let mut next_message_ix = prev_message_ix + 1;
2401            while let Some(next_message) = self.message_anchors.get(next_message_ix) {
2402                if next_message.start.is_valid(self.buffer.read(cx)) {
2403                    break;
2404                }
2405                next_message_ix += 1;
2406            }
2407
2408            let buffer = self.buffer.read(cx);
2409            let offset = self
2410                .message_anchors
2411                .get(next_message_ix)
2412                .map_or(buffer.len(), |message| {
2413                    buffer.clip_offset(message.start.to_offset(buffer) - 1, Bias::Left)
2414                });
2415            Some(self.insert_message_at_offset(offset, role, status, cx))
2416        } else {
2417            None
2418        }
2419    }
2420
2421    fn insert_message_at_offset(
2422        &mut self,
2423        offset: usize,
2424        role: Role,
2425        status: MessageStatus,
2426        cx: &mut Context<Self>,
2427    ) -> MessageAnchor {
2428        let start = self.buffer.update(cx, |buffer, cx| {
2429            buffer.edit([(offset..offset, "\n")], None, cx);
2430            buffer.anchor_before(offset + 1)
2431        });
2432
2433        let version = self.version.clone();
2434        let anchor = MessageAnchor {
2435            id: MessageId(self.next_timestamp()),
2436            start,
2437        };
2438        let metadata = MessageMetadata {
2439            role,
2440            status,
2441            timestamp: anchor.id.0,
2442            cache: None,
2443        };
2444        self.insert_message(anchor.clone(), metadata.clone(), cx);
2445        self.push_op(
2446            ContextOperation::InsertMessage {
2447                anchor: anchor.clone(),
2448                metadata,
2449                version,
2450            },
2451            cx,
2452        );
2453        anchor
2454    }
2455
2456    pub fn insert_content(&mut self, content: Content, cx: &mut Context<Self>) {
2457        let buffer = self.buffer.read(cx);
2458        let insertion_ix = match self
2459            .contents
2460            .binary_search_by(|probe| probe.cmp(&content, buffer))
2461        {
2462            Ok(ix) => {
2463                self.contents.remove(ix);
2464                ix
2465            }
2466            Err(ix) => ix,
2467        };
2468        self.contents.insert(insertion_ix, content);
2469        cx.emit(ContextEvent::MessagesEdited);
2470    }
2471
2472    pub fn contents<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Content> {
2473        let buffer = self.buffer.read(cx);
2474        self.contents
2475            .iter()
2476            .filter(|content| {
2477                let range = content.range();
2478                range.start.is_valid(buffer) && range.end.is_valid(buffer)
2479            })
2480            .cloned()
2481    }
2482
2483    pub fn split_message(
2484        &mut self,
2485        range: Range<usize>,
2486        cx: &mut Context<Self>,
2487    ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
2488        let start_message = self.message_for_offset(range.start, cx);
2489        let end_message = self.message_for_offset(range.end, cx);
2490        if let Some((start_message, end_message)) = start_message.zip(end_message) {
2491            // Prevent splitting when range spans multiple messages.
2492            if start_message.id != end_message.id {
2493                return (None, None);
2494            }
2495
2496            let message = start_message;
2497            let role = message.role;
2498            let mut edited_buffer = false;
2499
2500            let mut suffix_start = None;
2501
2502            // TODO: why did this start panicking?
2503            if range.start > message.offset_range.start
2504                && range.end < message.offset_range.end.saturating_sub(1)
2505            {
2506                if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
2507                    suffix_start = Some(range.end + 1);
2508                } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
2509                    suffix_start = Some(range.end);
2510                }
2511            }
2512
2513            let version = self.version.clone();
2514            let suffix = if let Some(suffix_start) = suffix_start {
2515                MessageAnchor {
2516                    id: MessageId(self.next_timestamp()),
2517                    start: self.buffer.read(cx).anchor_before(suffix_start),
2518                }
2519            } else {
2520                self.buffer.update(cx, |buffer, cx| {
2521                    buffer.edit([(range.end..range.end, "\n")], None, cx);
2522                });
2523                edited_buffer = true;
2524                MessageAnchor {
2525                    id: MessageId(self.next_timestamp()),
2526                    start: self.buffer.read(cx).anchor_before(range.end + 1),
2527                }
2528            };
2529
2530            let suffix_metadata = MessageMetadata {
2531                role,
2532                status: MessageStatus::Done,
2533                timestamp: suffix.id.0,
2534                cache: None,
2535            };
2536            self.insert_message(suffix.clone(), suffix_metadata.clone(), cx);
2537            self.push_op(
2538                ContextOperation::InsertMessage {
2539                    anchor: suffix.clone(),
2540                    metadata: suffix_metadata,
2541                    version,
2542                },
2543                cx,
2544            );
2545
2546            let new_messages =
2547                if range.start == range.end || range.start == message.offset_range.start {
2548                    (None, Some(suffix))
2549                } else {
2550                    let mut prefix_end = None;
2551                    if range.start > message.offset_range.start
2552                        && range.end < message.offset_range.end - 1
2553                    {
2554                        if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
2555                            prefix_end = Some(range.start + 1);
2556                        } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
2557                            == Some('\n')
2558                        {
2559                            prefix_end = Some(range.start);
2560                        }
2561                    }
2562
2563                    let version = self.version.clone();
2564                    let selection = if let Some(prefix_end) = prefix_end {
2565                        MessageAnchor {
2566                            id: MessageId(self.next_timestamp()),
2567                            start: self.buffer.read(cx).anchor_before(prefix_end),
2568                        }
2569                    } else {
2570                        self.buffer.update(cx, |buffer, cx| {
2571                            buffer.edit([(range.start..range.start, "\n")], None, cx)
2572                        });
2573                        edited_buffer = true;
2574                        MessageAnchor {
2575                            id: MessageId(self.next_timestamp()),
2576                            start: self.buffer.read(cx).anchor_before(range.end + 1),
2577                        }
2578                    };
2579
2580                    let selection_metadata = MessageMetadata {
2581                        role,
2582                        status: MessageStatus::Done,
2583                        timestamp: selection.id.0,
2584                        cache: None,
2585                    };
2586                    self.insert_message(selection.clone(), selection_metadata.clone(), cx);
2587                    self.push_op(
2588                        ContextOperation::InsertMessage {
2589                            anchor: selection.clone(),
2590                            metadata: selection_metadata,
2591                            version,
2592                        },
2593                        cx,
2594                    );
2595
2596                    (Some(selection), Some(suffix))
2597                };
2598
2599            if !edited_buffer {
2600                cx.emit(ContextEvent::MessagesEdited);
2601            }
2602            new_messages
2603        } else {
2604            (None, None)
2605        }
2606    }
2607
2608    fn insert_message(
2609        &mut self,
2610        new_anchor: MessageAnchor,
2611        new_metadata: MessageMetadata,
2612        cx: &mut Context<Self>,
2613    ) {
2614        cx.emit(ContextEvent::MessagesEdited);
2615
2616        self.messages_metadata.insert(new_anchor.id, new_metadata);
2617
2618        let buffer = self.buffer.read(cx);
2619        let insertion_ix = self
2620            .message_anchors
2621            .iter()
2622            .position(|anchor| {
2623                let comparison = new_anchor.start.cmp(&anchor.start, buffer);
2624                comparison.is_lt() || (comparison.is_eq() && new_anchor.id > anchor.id)
2625            })
2626            .unwrap_or(self.message_anchors.len());
2627        self.message_anchors.insert(insertion_ix, new_anchor);
2628    }
2629
2630    pub fn summarize(&mut self, mut replace_old: bool, cx: &mut Context<Self>) {
2631        let Some(model) = LanguageModelRegistry::read_global(cx).default_model() else {
2632            return;
2633        };
2634
2635        if replace_old || (self.message_anchors.len() >= 2 && self.summary.is_pending()) {
2636            if !model.provider.is_authenticated(cx) {
2637                return;
2638            }
2639
2640            let mut request = self.to_completion_request(Some(&model.model), cx);
2641            request.messages.push(LanguageModelRequestMessage {
2642                role: Role::User,
2643                content: vec![
2644                    "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:`"
2645                        .into(),
2646                ],
2647                cache: false,
2648            });
2649
2650            // If there is no summary, it is set with `done: false` so that "Loading Summary…" can
2651            // be displayed.
2652            match self.summary {
2653                ContextSummary::Pending | ContextSummary::Error => {
2654                    self.summary = ContextSummary::Content(ContextSummaryContent {
2655                        text: "".to_string(),
2656                        done: false,
2657                        timestamp: clock::Lamport::default(),
2658                    });
2659                    replace_old = true;
2660                }
2661                ContextSummary::Content(_) => {}
2662            }
2663
2664            self.summary_task = cx.spawn(async move |this, cx| {
2665                let result = async {
2666                    let stream = model.model.stream_completion_text(request, &cx);
2667                    let mut messages = stream.await?;
2668
2669                    let mut replaced = !replace_old;
2670                    while let Some(message) = messages.stream.next().await {
2671                        let text = message?;
2672                        let mut lines = text.lines();
2673                        this.update(cx, |this, cx| {
2674                            let version = this.version.clone();
2675                            let timestamp = this.next_timestamp();
2676                            let summary = this.summary.content_or_set_empty();
2677                            if !replaced && replace_old {
2678                                summary.text.clear();
2679                                replaced = true;
2680                            }
2681                            summary.text.extend(lines.next());
2682                            summary.timestamp = timestamp;
2683                            let operation = ContextOperation::UpdateSummary {
2684                                summary: summary.clone(),
2685                                version,
2686                            };
2687                            this.push_op(operation, cx);
2688                            cx.emit(ContextEvent::SummaryChanged);
2689                            cx.emit(ContextEvent::SummaryGenerated);
2690                        })?;
2691
2692                        // Stop if the LLM generated multiple lines.
2693                        if lines.next().is_some() {
2694                            break;
2695                        }
2696                    }
2697
2698                    this.read_with(cx, |this, _cx| {
2699                        if let Some(summary) = this.summary.content() {
2700                            if summary.text.is_empty() {
2701                                bail!("Model generated an empty summary");
2702                            }
2703                        }
2704                        Ok(())
2705                    })??;
2706
2707                    this.update(cx, |this, cx| {
2708                        let version = this.version.clone();
2709                        let timestamp = this.next_timestamp();
2710                        if let Some(summary) = this.summary.content_as_mut() {
2711                            summary.done = true;
2712                            summary.timestamp = timestamp;
2713                            let operation = ContextOperation::UpdateSummary {
2714                                summary: summary.clone(),
2715                                version,
2716                            };
2717                            this.push_op(operation, cx);
2718                            cx.emit(ContextEvent::SummaryChanged);
2719                            cx.emit(ContextEvent::SummaryGenerated);
2720                        }
2721                    })?;
2722
2723                    anyhow::Ok(())
2724                }
2725                .await;
2726
2727                if let Err(err) = result {
2728                    this.update(cx, |this, cx| {
2729                        this.summary = ContextSummary::Error;
2730                        cx.emit(ContextEvent::SummaryChanged);
2731                    })
2732                    .log_err();
2733                    log::error!("Error generating context summary: {}", err);
2734                }
2735
2736                Some(())
2737            });
2738        }
2739    }
2740
2741    fn message_for_offset(&self, offset: usize, cx: &App) -> Option<Message> {
2742        self.messages_for_offsets([offset], cx).pop()
2743    }
2744
2745    pub fn messages_for_offsets(
2746        &self,
2747        offsets: impl IntoIterator<Item = usize>,
2748        cx: &App,
2749    ) -> Vec<Message> {
2750        let mut result = Vec::new();
2751
2752        let mut messages = self.messages(cx).peekable();
2753        let mut offsets = offsets.into_iter().peekable();
2754        let mut current_message = messages.next();
2755        while let Some(offset) = offsets.next() {
2756            // Locate the message that contains the offset.
2757            while current_message.as_ref().map_or(false, |message| {
2758                !message.offset_range.contains(&offset) && messages.peek().is_some()
2759            }) {
2760                current_message = messages.next();
2761            }
2762            let Some(message) = current_message.as_ref() else {
2763                break;
2764            };
2765
2766            // Skip offsets that are in the same message.
2767            while offsets.peek().map_or(false, |offset| {
2768                message.offset_range.contains(offset) || messages.peek().is_none()
2769            }) {
2770                offsets.next();
2771            }
2772
2773            result.push(message.clone());
2774        }
2775        result
2776    }
2777
2778    fn messages_from_anchors<'a>(
2779        &'a self,
2780        message_anchors: impl Iterator<Item = &'a MessageAnchor> + 'a,
2781        cx: &'a App,
2782    ) -> impl 'a + Iterator<Item = Message> {
2783        let buffer = self.buffer.read(cx);
2784
2785        Self::messages_from_iters(buffer, &self.messages_metadata, message_anchors.enumerate())
2786    }
2787
2788    pub fn messages<'a>(&'a self, cx: &'a App) -> impl 'a + Iterator<Item = Message> {
2789        self.messages_from_anchors(self.message_anchors.iter(), cx)
2790    }
2791
2792    pub fn messages_from_iters<'a>(
2793        buffer: &'a Buffer,
2794        metadata: &'a HashMap<MessageId, MessageMetadata>,
2795        messages: impl Iterator<Item = (usize, &'a MessageAnchor)> + 'a,
2796    ) -> impl 'a + Iterator<Item = Message> {
2797        let mut messages = messages.peekable();
2798
2799        iter::from_fn(move || {
2800            if let Some((start_ix, message_anchor)) = messages.next() {
2801                let metadata = metadata.get(&message_anchor.id)?;
2802
2803                let message_start = message_anchor.start.to_offset(buffer);
2804                let mut message_end = None;
2805                let mut end_ix = start_ix;
2806                while let Some((_, next_message)) = messages.peek() {
2807                    if next_message.start.is_valid(buffer) {
2808                        message_end = Some(next_message.start);
2809                        break;
2810                    } else {
2811                        end_ix += 1;
2812                        messages.next();
2813                    }
2814                }
2815                let message_end_anchor = message_end.unwrap_or(language::Anchor::MAX);
2816                let message_end = message_end_anchor.to_offset(buffer);
2817
2818                return Some(Message {
2819                    index_range: start_ix..end_ix,
2820                    offset_range: message_start..message_end,
2821                    anchor_range: message_anchor.start..message_end_anchor,
2822                    id: message_anchor.id,
2823                    role: metadata.role,
2824                    status: metadata.status.clone(),
2825                    cache: metadata.cache.clone(),
2826                });
2827            }
2828            None
2829        })
2830    }
2831
2832    pub fn save(
2833        &mut self,
2834        debounce: Option<Duration>,
2835        fs: Arc<dyn Fs>,
2836        cx: &mut Context<AssistantContext>,
2837    ) {
2838        if self.replica_id() != ReplicaId::default() {
2839            // Prevent saving a remote context for now.
2840            return;
2841        }
2842
2843        self.pending_save = cx.spawn(async move |this, cx| {
2844            if let Some(debounce) = debounce {
2845                cx.background_executor().timer(debounce).await;
2846            }
2847
2848            let (old_path, summary) = this.read_with(cx, |this, _| {
2849                let path = this.path.clone();
2850                let summary = if let Some(summary) = this.summary.content() {
2851                    if summary.done {
2852                        Some(summary.text.clone())
2853                    } else {
2854                        None
2855                    }
2856                } else {
2857                    None
2858                };
2859                (path, summary)
2860            })?;
2861
2862            if let Some(summary) = summary {
2863                let context = this.read_with(cx, |this, cx| this.serialize(cx))?;
2864                let mut discriminant = 1;
2865                let mut new_path;
2866                loop {
2867                    new_path = contexts_dir().join(&format!(
2868                        "{} - {}.zed.json",
2869                        summary.trim(),
2870                        discriminant
2871                    ));
2872                    if fs.is_file(&new_path).await {
2873                        discriminant += 1;
2874                    } else {
2875                        break;
2876                    }
2877                }
2878
2879                fs.create_dir(contexts_dir().as_ref()).await?;
2880                fs.atomic_write(new_path.clone(), serde_json::to_string(&context).unwrap())
2881                    .await?;
2882                if let Some(old_path) = old_path {
2883                    if new_path.as_path() != old_path.as_ref() {
2884                        fs.remove_file(
2885                            &old_path,
2886                            RemoveOptions {
2887                                recursive: false,
2888                                ignore_if_not_exists: true,
2889                            },
2890                        )
2891                        .await?;
2892                    }
2893                }
2894
2895                this.update(cx, |this, _| this.path = Some(new_path.into()))?;
2896            }
2897
2898            Ok(())
2899        });
2900    }
2901
2902    pub fn set_custom_summary(&mut self, custom_summary: String, cx: &mut Context<Self>) {
2903        let timestamp = self.next_timestamp();
2904        let summary = self.summary.content_or_set_empty();
2905        summary.timestamp = timestamp;
2906        summary.done = true;
2907        summary.text = custom_summary;
2908        cx.emit(ContextEvent::SummaryChanged);
2909    }
2910}
2911
2912#[derive(Debug, Default)]
2913pub struct ContextVersion {
2914    context: clock::Global,
2915    buffer: clock::Global,
2916}
2917
2918impl ContextVersion {
2919    pub fn from_proto(proto: &proto::ContextVersion) -> Self {
2920        Self {
2921            context: language::proto::deserialize_version(&proto.context_version),
2922            buffer: language::proto::deserialize_version(&proto.buffer_version),
2923        }
2924    }
2925
2926    pub fn to_proto(&self, context_id: ContextId) -> proto::ContextVersion {
2927        proto::ContextVersion {
2928            context_id: context_id.to_proto(),
2929            context_version: language::proto::serialize_version(&self.context),
2930            buffer_version: language::proto::serialize_version(&self.buffer),
2931        }
2932    }
2933}
2934
2935#[derive(Debug, Clone)]
2936pub struct ParsedSlashCommand {
2937    pub name: String,
2938    pub arguments: SmallVec<[String; 3]>,
2939    pub status: PendingSlashCommandStatus,
2940    pub source_range: Range<language::Anchor>,
2941}
2942
2943#[derive(Debug)]
2944pub struct InvokedSlashCommand {
2945    pub name: SharedString,
2946    pub range: Range<language::Anchor>,
2947    pub run_commands_in_ranges: Vec<Range<language::Anchor>>,
2948    pub status: InvokedSlashCommandStatus,
2949    pub transaction: Option<language::TransactionId>,
2950    timestamp: clock::Lamport,
2951}
2952
2953#[derive(Debug)]
2954pub enum InvokedSlashCommandStatus {
2955    Running(Task<()>),
2956    Error(SharedString),
2957    Finished,
2958}
2959
2960#[derive(Debug, Clone)]
2961pub enum PendingSlashCommandStatus {
2962    Idle,
2963    Running { _task: Shared<Task<()>> },
2964    Error(String),
2965}
2966
2967#[derive(Debug, Clone)]
2968pub struct PendingToolUse {
2969    pub id: LanguageModelToolUseId,
2970    pub name: String,
2971    pub input: serde_json::Value,
2972    pub status: PendingToolUseStatus,
2973    pub source_range: Range<language::Anchor>,
2974}
2975
2976#[derive(Debug, Clone)]
2977pub enum PendingToolUseStatus {
2978    Idle,
2979    Running { _task: Shared<Task<()>> },
2980    Error(String),
2981}
2982
2983impl PendingToolUseStatus {
2984    pub fn is_idle(&self) -> bool {
2985        matches!(self, PendingToolUseStatus::Idle)
2986    }
2987}
2988
2989#[derive(Serialize, Deserialize)]
2990pub struct SavedMessage {
2991    pub id: MessageId,
2992    pub start: usize,
2993    pub metadata: MessageMetadata,
2994}
2995
2996#[derive(Serialize, Deserialize)]
2997pub struct SavedContext {
2998    pub id: Option<ContextId>,
2999    pub zed: String,
3000    pub version: String,
3001    pub text: String,
3002    pub messages: Vec<SavedMessage>,
3003    pub summary: String,
3004    pub slash_command_output_sections:
3005        Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3006    #[serde(default)]
3007    pub thought_process_output_sections: Vec<ThoughtProcessOutputSection<usize>>,
3008}
3009
3010impl SavedContext {
3011    pub const VERSION: &'static str = "0.4.0";
3012
3013    pub fn from_json(json: &str) -> Result<Self> {
3014        let saved_context_json = serde_json::from_str::<serde_json::Value>(json)?;
3015        match saved_context_json
3016            .get("version")
3017            .context("version not found")?
3018        {
3019            serde_json::Value::String(version) => match version.as_str() {
3020                SavedContext::VERSION => {
3021                    Ok(serde_json::from_value::<SavedContext>(saved_context_json)?)
3022                }
3023                SavedContextV0_3_0::VERSION => {
3024                    let saved_context =
3025                        serde_json::from_value::<SavedContextV0_3_0>(saved_context_json)?;
3026                    Ok(saved_context.upgrade())
3027                }
3028                SavedContextV0_2_0::VERSION => {
3029                    let saved_context =
3030                        serde_json::from_value::<SavedContextV0_2_0>(saved_context_json)?;
3031                    Ok(saved_context.upgrade())
3032                }
3033                SavedContextV0_1_0::VERSION => {
3034                    let saved_context =
3035                        serde_json::from_value::<SavedContextV0_1_0>(saved_context_json)?;
3036                    Ok(saved_context.upgrade())
3037                }
3038                _ => anyhow::bail!("unrecognized saved context version: {version:?}"),
3039            },
3040            _ => anyhow::bail!("version not found on saved context"),
3041        }
3042    }
3043
3044    fn into_ops(
3045        self,
3046        buffer: &Entity<Buffer>,
3047        cx: &mut Context<AssistantContext>,
3048    ) -> Vec<ContextOperation> {
3049        let mut operations = Vec::new();
3050        let mut version = clock::Global::new();
3051        let mut next_timestamp = clock::Lamport::new(ReplicaId::default());
3052
3053        let mut first_message_metadata = None;
3054        for message in self.messages {
3055            if message.id == MessageId(clock::Lamport::default()) {
3056                first_message_metadata = Some(message.metadata);
3057            } else {
3058                operations.push(ContextOperation::InsertMessage {
3059                    anchor: MessageAnchor {
3060                        id: message.id,
3061                        start: buffer.read(cx).anchor_before(message.start),
3062                    },
3063                    metadata: MessageMetadata {
3064                        role: message.metadata.role,
3065                        status: message.metadata.status,
3066                        timestamp: message.metadata.timestamp,
3067                        cache: None,
3068                    },
3069                    version: version.clone(),
3070                });
3071                version.observe(message.id.0);
3072                next_timestamp.observe(message.id.0);
3073            }
3074        }
3075
3076        if let Some(metadata) = first_message_metadata {
3077            let timestamp = next_timestamp.tick();
3078            operations.push(ContextOperation::UpdateMessage {
3079                message_id: MessageId(clock::Lamport::default()),
3080                metadata: MessageMetadata {
3081                    role: metadata.role,
3082                    status: metadata.status,
3083                    timestamp,
3084                    cache: None,
3085                },
3086                version: version.clone(),
3087            });
3088            version.observe(timestamp);
3089        }
3090
3091        let buffer = buffer.read(cx);
3092        for section in self.slash_command_output_sections {
3093            let timestamp = next_timestamp.tick();
3094            operations.push(ContextOperation::SlashCommandOutputSectionAdded {
3095                timestamp,
3096                section: SlashCommandOutputSection {
3097                    range: buffer.anchor_after(section.range.start)
3098                        ..buffer.anchor_before(section.range.end),
3099                    icon: section.icon,
3100                    label: section.label,
3101                    metadata: section.metadata,
3102                },
3103                version: version.clone(),
3104            });
3105
3106            version.observe(timestamp);
3107        }
3108
3109        for section in self.thought_process_output_sections {
3110            let timestamp = next_timestamp.tick();
3111            operations.push(ContextOperation::ThoughtProcessOutputSectionAdded {
3112                timestamp,
3113                section: ThoughtProcessOutputSection {
3114                    range: buffer.anchor_after(section.range.start)
3115                        ..buffer.anchor_before(section.range.end),
3116                },
3117                version: version.clone(),
3118            });
3119
3120            version.observe(timestamp);
3121        }
3122
3123        let timestamp = next_timestamp.tick();
3124        operations.push(ContextOperation::UpdateSummary {
3125            summary: ContextSummaryContent {
3126                text: self.summary,
3127                done: true,
3128                timestamp,
3129            },
3130            version: version.clone(),
3131        });
3132        version.observe(timestamp);
3133
3134        operations
3135    }
3136}
3137
3138#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
3139struct SavedMessageIdPreV0_4_0(usize);
3140
3141#[derive(Serialize, Deserialize)]
3142struct SavedMessagePreV0_4_0 {
3143    id: SavedMessageIdPreV0_4_0,
3144    start: usize,
3145}
3146
3147#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
3148struct SavedMessageMetadataPreV0_4_0 {
3149    role: Role,
3150    status: MessageStatus,
3151}
3152
3153#[derive(Serialize, Deserialize)]
3154struct SavedContextV0_3_0 {
3155    id: Option<ContextId>,
3156    zed: String,
3157    version: String,
3158    text: String,
3159    messages: Vec<SavedMessagePreV0_4_0>,
3160    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3161    summary: String,
3162    slash_command_output_sections: Vec<assistant_slash_command::SlashCommandOutputSection<usize>>,
3163}
3164
3165impl SavedContextV0_3_0 {
3166    const VERSION: &'static str = "0.3.0";
3167
3168    fn upgrade(self) -> SavedContext {
3169        SavedContext {
3170            id: self.id,
3171            zed: self.zed,
3172            version: SavedContext::VERSION.into(),
3173            text: self.text,
3174            messages: self
3175                .messages
3176                .into_iter()
3177                .filter_map(|message| {
3178                    let metadata = self.message_metadata.get(&message.id)?;
3179                    let timestamp = clock::Lamport {
3180                        replica_id: ReplicaId::default(),
3181                        value: message.id.0 as u32,
3182                    };
3183                    Some(SavedMessage {
3184                        id: MessageId(timestamp),
3185                        start: message.start,
3186                        metadata: MessageMetadata {
3187                            role: metadata.role,
3188                            status: metadata.status.clone(),
3189                            timestamp,
3190                            cache: None,
3191                        },
3192                    })
3193                })
3194                .collect(),
3195            summary: self.summary,
3196            slash_command_output_sections: self.slash_command_output_sections,
3197            thought_process_output_sections: Vec::new(),
3198        }
3199    }
3200}
3201
3202#[derive(Serialize, Deserialize)]
3203struct SavedContextV0_2_0 {
3204    id: Option<ContextId>,
3205    zed: String,
3206    version: String,
3207    text: String,
3208    messages: Vec<SavedMessagePreV0_4_0>,
3209    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3210    summary: String,
3211}
3212
3213impl SavedContextV0_2_0 {
3214    const VERSION: &'static str = "0.2.0";
3215
3216    fn upgrade(self) -> SavedContext {
3217        SavedContextV0_3_0 {
3218            id: self.id,
3219            zed: self.zed,
3220            version: SavedContextV0_3_0::VERSION.to_string(),
3221            text: self.text,
3222            messages: self.messages,
3223            message_metadata: self.message_metadata,
3224            summary: self.summary,
3225            slash_command_output_sections: Vec::new(),
3226        }
3227        .upgrade()
3228    }
3229}
3230
3231#[derive(Serialize, Deserialize)]
3232struct SavedContextV0_1_0 {
3233    id: Option<ContextId>,
3234    zed: String,
3235    version: String,
3236    text: String,
3237    messages: Vec<SavedMessagePreV0_4_0>,
3238    message_metadata: HashMap<SavedMessageIdPreV0_4_0, SavedMessageMetadataPreV0_4_0>,
3239    summary: String,
3240    api_url: Option<String>,
3241    model: OpenAiModel,
3242}
3243
3244impl SavedContextV0_1_0 {
3245    const VERSION: &'static str = "0.1.0";
3246
3247    fn upgrade(self) -> SavedContext {
3248        SavedContextV0_2_0 {
3249            id: self.id,
3250            zed: self.zed,
3251            version: SavedContextV0_2_0::VERSION.to_string(),
3252            text: self.text,
3253            messages: self.messages,
3254            message_metadata: self.message_metadata,
3255            summary: self.summary,
3256        }
3257        .upgrade()
3258    }
3259}
3260
3261#[derive(Debug, Clone)]
3262pub struct SavedContextMetadata {
3263    pub title: String,
3264    pub path: Arc<Path>,
3265    pub mtime: chrono::DateTime<chrono::Local>,
3266}