assistant_context.rs

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