thread.rs

   1use std::fmt::Write as _;
   2use std::io::Write;
   3use std::ops::Range;
   4use std::sync::Arc;
   5
   6use anyhow::{Context as _, Result, anyhow};
   7use assistant_settings::AssistantSettings;
   8use assistant_tool::{ActionLog, Tool, ToolWorkingSet};
   9use chrono::{DateTime, Utc};
  10use collections::{BTreeMap, HashMap};
  11use fs::Fs;
  12use futures::future::Shared;
  13use futures::{FutureExt, StreamExt as _};
  14use git::repository::DiffType;
  15use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
  16use language_model::{
  17    ConfiguredModel, LanguageModel, LanguageModelCompletionEvent, LanguageModelRegistry,
  18    LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool,
  19    LanguageModelToolResult, LanguageModelToolUseId, MaxMonthlySpendReachedError, MessageContent,
  20    PaymentRequiredError, Role, StopReason, TokenUsage,
  21};
  22use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState};
  23use project::{Project, Worktree};
  24use prompt_store::{
  25    AssistantSystemPromptContext, PromptBuilder, RulesFile, WorktreeInfoForSystemPrompt,
  26};
  27use schemars::JsonSchema;
  28use serde::{Deserialize, Serialize};
  29use settings::Settings;
  30use util::{ResultExt as _, TryFutureExt as _, maybe, post_inc};
  31use uuid::Uuid;
  32
  33use crate::context::{AssistantContext, ContextId, format_context_as_string};
  34use crate::thread_store::{
  35    SerializedMessage, SerializedMessageSegment, SerializedThread, SerializedToolResult,
  36    SerializedToolUse,
  37};
  38use crate::tool_use::{PendingToolUse, ToolUse, ToolUseState, USING_TOOL_MARKER};
  39
  40#[derive(Debug, Clone, Copy)]
  41pub enum RequestKind {
  42    Chat,
  43    /// Used when summarizing a thread.
  44    Summarize,
  45}
  46
  47#[derive(
  48    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
  49)]
  50pub struct ThreadId(Arc<str>);
  51
  52impl ThreadId {
  53    pub fn new() -> Self {
  54        Self(Uuid::new_v4().to_string().into())
  55    }
  56}
  57
  58impl std::fmt::Display for ThreadId {
  59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  60        write!(f, "{}", self.0)
  61    }
  62}
  63
  64impl From<&str> for ThreadId {
  65    fn from(value: &str) -> Self {
  66        Self(value.into())
  67    }
  68}
  69
  70#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
  71pub struct MessageId(pub(crate) usize);
  72
  73impl MessageId {
  74    fn post_inc(&mut self) -> Self {
  75        Self(post_inc(&mut self.0))
  76    }
  77}
  78
  79/// A message in a [`Thread`].
  80#[derive(Debug, Clone)]
  81pub struct Message {
  82    pub id: MessageId,
  83    pub role: Role,
  84    pub segments: Vec<MessageSegment>,
  85    pub context: String,
  86}
  87
  88impl Message {
  89    /// Returns whether the message contains any meaningful text that should be displayed
  90    /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
  91    pub fn should_display_content(&self) -> bool {
  92        self.segments.iter().all(|segment| segment.should_display())
  93    }
  94
  95    pub fn push_thinking(&mut self, text: &str) {
  96        if let Some(MessageSegment::Thinking(segment)) = self.segments.last_mut() {
  97            segment.push_str(text);
  98        } else {
  99            self.segments
 100                .push(MessageSegment::Thinking(text.to_string()));
 101        }
 102    }
 103
 104    pub fn push_text(&mut self, text: &str) {
 105        if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
 106            segment.push_str(text);
 107        } else {
 108            self.segments.push(MessageSegment::Text(text.to_string()));
 109        }
 110    }
 111
 112    pub fn to_string(&self) -> String {
 113        let mut result = String::new();
 114
 115        if !self.context.is_empty() {
 116            result.push_str(&self.context);
 117        }
 118
 119        for segment in &self.segments {
 120            match segment {
 121                MessageSegment::Text(text) => result.push_str(text),
 122                MessageSegment::Thinking(text) => {
 123                    result.push_str("<think>");
 124                    result.push_str(text);
 125                    result.push_str("</think>");
 126                }
 127            }
 128        }
 129
 130        result
 131    }
 132}
 133
 134#[derive(Debug, Clone, PartialEq, Eq)]
 135pub enum MessageSegment {
 136    Text(String),
 137    Thinking(String),
 138}
 139
 140impl MessageSegment {
 141    pub fn text_mut(&mut self) -> &mut String {
 142        match self {
 143            Self::Text(text) => text,
 144            Self::Thinking(text) => text,
 145        }
 146    }
 147
 148    pub fn should_display(&self) -> bool {
 149        // We add USING_TOOL_MARKER when making a request that includes tool uses
 150        // without non-whitespace text around them, and this can cause the model
 151        // to mimic the pattern, so we consider those segments not displayable.
 152        match self {
 153            Self::Text(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
 154            Self::Thinking(text) => text.is_empty() || text.trim() == USING_TOOL_MARKER,
 155        }
 156    }
 157}
 158
 159#[derive(Debug, Clone, Serialize, Deserialize)]
 160pub struct ProjectSnapshot {
 161    pub worktree_snapshots: Vec<WorktreeSnapshot>,
 162    pub unsaved_buffer_paths: Vec<String>,
 163    pub timestamp: DateTime<Utc>,
 164}
 165
 166#[derive(Debug, Clone, Serialize, Deserialize)]
 167pub struct WorktreeSnapshot {
 168    pub worktree_path: String,
 169    pub git_state: Option<GitState>,
 170}
 171
 172#[derive(Debug, Clone, Serialize, Deserialize)]
 173pub struct GitState {
 174    pub remote_url: Option<String>,
 175    pub head_sha: Option<String>,
 176    pub current_branch: Option<String>,
 177    pub diff: Option<String>,
 178}
 179
 180#[derive(Clone)]
 181pub struct ThreadCheckpoint {
 182    message_id: MessageId,
 183    git_checkpoint: GitStoreCheckpoint,
 184}
 185
 186#[derive(Copy, Clone, Debug)]
 187pub enum ThreadFeedback {
 188    Positive,
 189    Negative,
 190}
 191
 192pub enum LastRestoreCheckpoint {
 193    Pending {
 194        message_id: MessageId,
 195    },
 196    Error {
 197        message_id: MessageId,
 198        error: String,
 199    },
 200}
 201
 202impl LastRestoreCheckpoint {
 203    pub fn message_id(&self) -> MessageId {
 204        match self {
 205            LastRestoreCheckpoint::Pending { message_id } => *message_id,
 206            LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
 207        }
 208    }
 209}
 210
 211#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 212pub enum DetailedSummaryState {
 213    #[default]
 214    NotGenerated,
 215    Generating {
 216        message_id: MessageId,
 217    },
 218    Generated {
 219        text: SharedString,
 220        message_id: MessageId,
 221    },
 222}
 223
 224#[derive(Default)]
 225pub struct TotalTokenUsage {
 226    pub total: usize,
 227    pub max: usize,
 228    pub ratio: TokenUsageRatio,
 229}
 230
 231#[derive(Default, PartialEq, Eq)]
 232pub enum TokenUsageRatio {
 233    #[default]
 234    Normal,
 235    Warning,
 236    Exceeded,
 237}
 238
 239/// A thread of conversation with the LLM.
 240pub struct Thread {
 241    id: ThreadId,
 242    updated_at: DateTime<Utc>,
 243    summary: Option<SharedString>,
 244    pending_summary: Task<Option<()>>,
 245    detailed_summary_state: DetailedSummaryState,
 246    messages: Vec<Message>,
 247    next_message_id: MessageId,
 248    context: BTreeMap<ContextId, AssistantContext>,
 249    context_by_message: HashMap<MessageId, Vec<ContextId>>,
 250    system_prompt_context: Option<AssistantSystemPromptContext>,
 251    checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
 252    completion_count: usize,
 253    pending_completions: Vec<PendingCompletion>,
 254    project: Entity<Project>,
 255    prompt_builder: Arc<PromptBuilder>,
 256    tools: Arc<ToolWorkingSet>,
 257    tool_use: ToolUseState,
 258    action_log: Entity<ActionLog>,
 259    last_restore_checkpoint: Option<LastRestoreCheckpoint>,
 260    pending_checkpoint: Option<ThreadCheckpoint>,
 261    initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
 262    cumulative_token_usage: TokenUsage,
 263    feedback: Option<ThreadFeedback>,
 264}
 265
 266impl Thread {
 267    pub fn new(
 268        project: Entity<Project>,
 269        tools: Arc<ToolWorkingSet>,
 270        prompt_builder: Arc<PromptBuilder>,
 271        cx: &mut Context<Self>,
 272    ) -> Self {
 273        Self {
 274            id: ThreadId::new(),
 275            updated_at: Utc::now(),
 276            summary: None,
 277            pending_summary: Task::ready(None),
 278            detailed_summary_state: DetailedSummaryState::NotGenerated,
 279            messages: Vec::new(),
 280            next_message_id: MessageId(0),
 281            context: BTreeMap::default(),
 282            context_by_message: HashMap::default(),
 283            system_prompt_context: None,
 284            checkpoints_by_message: HashMap::default(),
 285            completion_count: 0,
 286            pending_completions: Vec::new(),
 287            project: project.clone(),
 288            prompt_builder,
 289            tools: tools.clone(),
 290            last_restore_checkpoint: None,
 291            pending_checkpoint: None,
 292            tool_use: ToolUseState::new(tools.clone()),
 293            action_log: cx.new(|_| ActionLog::new(project.clone())),
 294            initial_project_snapshot: {
 295                let project_snapshot = Self::project_snapshot(project, cx);
 296                cx.foreground_executor()
 297                    .spawn(async move { Some(project_snapshot.await) })
 298                    .shared()
 299            },
 300            cumulative_token_usage: TokenUsage::default(),
 301            feedback: None,
 302        }
 303    }
 304
 305    pub fn deserialize(
 306        id: ThreadId,
 307        serialized: SerializedThread,
 308        project: Entity<Project>,
 309        tools: Arc<ToolWorkingSet>,
 310        prompt_builder: Arc<PromptBuilder>,
 311        cx: &mut Context<Self>,
 312    ) -> Self {
 313        let next_message_id = MessageId(
 314            serialized
 315                .messages
 316                .last()
 317                .map(|message| message.id.0 + 1)
 318                .unwrap_or(0),
 319        );
 320        let tool_use =
 321            ToolUseState::from_serialized_messages(tools.clone(), &serialized.messages, |_| true);
 322
 323        Self {
 324            id,
 325            updated_at: serialized.updated_at,
 326            summary: Some(serialized.summary),
 327            pending_summary: Task::ready(None),
 328            detailed_summary_state: serialized.detailed_summary_state,
 329            messages: serialized
 330                .messages
 331                .into_iter()
 332                .map(|message| Message {
 333                    id: message.id,
 334                    role: message.role,
 335                    segments: message
 336                        .segments
 337                        .into_iter()
 338                        .map(|segment| match segment {
 339                            SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
 340                            SerializedMessageSegment::Thinking { text } => {
 341                                MessageSegment::Thinking(text)
 342                            }
 343                        })
 344                        .collect(),
 345                    context: message.context,
 346                })
 347                .collect(),
 348            next_message_id,
 349            context: BTreeMap::default(),
 350            context_by_message: HashMap::default(),
 351            system_prompt_context: None,
 352            checkpoints_by_message: HashMap::default(),
 353            completion_count: 0,
 354            pending_completions: Vec::new(),
 355            last_restore_checkpoint: None,
 356            pending_checkpoint: None,
 357            project: project.clone(),
 358            prompt_builder,
 359            tools,
 360            tool_use,
 361            action_log: cx.new(|_| ActionLog::new(project)),
 362            initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
 363            cumulative_token_usage: serialized.cumulative_token_usage,
 364            feedback: None,
 365        }
 366    }
 367
 368    pub fn id(&self) -> &ThreadId {
 369        &self.id
 370    }
 371
 372    pub fn is_empty(&self) -> bool {
 373        self.messages.is_empty()
 374    }
 375
 376    pub fn updated_at(&self) -> DateTime<Utc> {
 377        self.updated_at
 378    }
 379
 380    pub fn touch_updated_at(&mut self) {
 381        self.updated_at = Utc::now();
 382    }
 383
 384    pub fn summary(&self) -> Option<SharedString> {
 385        self.summary.clone()
 386    }
 387
 388    pub const DEFAULT_SUMMARY: SharedString = SharedString::new_static("New Thread");
 389
 390    pub fn summary_or_default(&self) -> SharedString {
 391        self.summary.clone().unwrap_or(Self::DEFAULT_SUMMARY)
 392    }
 393
 394    pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
 395        let Some(current_summary) = &self.summary else {
 396            // Don't allow setting summary until generated
 397            return;
 398        };
 399
 400        let mut new_summary = new_summary.into();
 401
 402        if new_summary.is_empty() {
 403            new_summary = Self::DEFAULT_SUMMARY;
 404        }
 405
 406        if current_summary != &new_summary {
 407            self.summary = Some(new_summary);
 408            cx.emit(ThreadEvent::SummaryChanged);
 409        }
 410    }
 411
 412    pub fn latest_detailed_summary_or_text(&self) -> SharedString {
 413        self.latest_detailed_summary()
 414            .unwrap_or_else(|| self.text().into())
 415    }
 416
 417    fn latest_detailed_summary(&self) -> Option<SharedString> {
 418        if let DetailedSummaryState::Generated { text, .. } = &self.detailed_summary_state {
 419            Some(text.clone())
 420        } else {
 421            None
 422        }
 423    }
 424
 425    pub fn message(&self, id: MessageId) -> Option<&Message> {
 426        self.messages.iter().find(|message| message.id == id)
 427    }
 428
 429    pub fn messages(&self) -> impl Iterator<Item = &Message> {
 430        self.messages.iter()
 431    }
 432
 433    pub fn is_generating(&self) -> bool {
 434        !self.pending_completions.is_empty() || !self.all_tools_finished()
 435    }
 436
 437    pub fn tools(&self) -> &Arc<ToolWorkingSet> {
 438        &self.tools
 439    }
 440
 441    pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
 442        self.tool_use
 443            .pending_tool_uses()
 444            .into_iter()
 445            .find(|tool_use| &tool_use.id == id)
 446    }
 447
 448    pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
 449        self.tool_use
 450            .pending_tool_uses()
 451            .into_iter()
 452            .filter(|tool_use| tool_use.status.needs_confirmation())
 453    }
 454
 455    pub fn has_pending_tool_uses(&self) -> bool {
 456        !self.tool_use.pending_tool_uses().is_empty()
 457    }
 458
 459    pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
 460        self.checkpoints_by_message.get(&id).cloned()
 461    }
 462
 463    pub fn restore_checkpoint(
 464        &mut self,
 465        checkpoint: ThreadCheckpoint,
 466        cx: &mut Context<Self>,
 467    ) -> Task<Result<()>> {
 468        self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
 469            message_id: checkpoint.message_id,
 470        });
 471        cx.emit(ThreadEvent::CheckpointChanged);
 472        cx.notify();
 473
 474        let git_store = self.project().read(cx).git_store().clone();
 475        let restore = git_store.update(cx, |git_store, cx| {
 476            git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
 477        });
 478
 479        cx.spawn(async move |this, cx| {
 480            let result = restore.await;
 481            this.update(cx, |this, cx| {
 482                if let Err(err) = result.as_ref() {
 483                    this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
 484                        message_id: checkpoint.message_id,
 485                        error: err.to_string(),
 486                    });
 487                } else {
 488                    this.truncate(checkpoint.message_id, cx);
 489                    this.last_restore_checkpoint = None;
 490                }
 491                this.pending_checkpoint = None;
 492                cx.emit(ThreadEvent::CheckpointChanged);
 493                cx.notify();
 494            })?;
 495            result
 496        })
 497    }
 498
 499    fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
 500        let pending_checkpoint = if self.is_generating() {
 501            return;
 502        } else if let Some(checkpoint) = self.pending_checkpoint.take() {
 503            checkpoint
 504        } else {
 505            return;
 506        };
 507
 508        let git_store = self.project.read(cx).git_store().clone();
 509        let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
 510        cx.spawn(async move |this, cx| match final_checkpoint.await {
 511            Ok(final_checkpoint) => {
 512                let equal = git_store
 513                    .update(cx, |store, cx| {
 514                        store.compare_checkpoints(
 515                            pending_checkpoint.git_checkpoint.clone(),
 516                            final_checkpoint.clone(),
 517                            cx,
 518                        )
 519                    })?
 520                    .await
 521                    .unwrap_or(false);
 522
 523                if equal {
 524                    git_store
 525                        .update(cx, |store, cx| {
 526                            store.delete_checkpoint(pending_checkpoint.git_checkpoint, cx)
 527                        })?
 528                        .detach();
 529                } else {
 530                    this.update(cx, |this, cx| {
 531                        this.insert_checkpoint(pending_checkpoint, cx)
 532                    })?;
 533                }
 534
 535                git_store
 536                    .update(cx, |store, cx| {
 537                        store.delete_checkpoint(final_checkpoint, cx)
 538                    })?
 539                    .detach();
 540
 541                Ok(())
 542            }
 543            Err(_) => this.update(cx, |this, cx| {
 544                this.insert_checkpoint(pending_checkpoint, cx)
 545            }),
 546        })
 547        .detach();
 548    }
 549
 550    fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
 551        self.checkpoints_by_message
 552            .insert(checkpoint.message_id, checkpoint);
 553        cx.emit(ThreadEvent::CheckpointChanged);
 554        cx.notify();
 555    }
 556
 557    pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
 558        self.last_restore_checkpoint.as_ref()
 559    }
 560
 561    pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
 562        let Some(message_ix) = self
 563            .messages
 564            .iter()
 565            .rposition(|message| message.id == message_id)
 566        else {
 567            return;
 568        };
 569        for deleted_message in self.messages.drain(message_ix..) {
 570            self.context_by_message.remove(&deleted_message.id);
 571            self.checkpoints_by_message.remove(&deleted_message.id);
 572        }
 573        cx.notify();
 574    }
 575
 576    pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AssistantContext> {
 577        self.context_by_message
 578            .get(&id)
 579            .into_iter()
 580            .flat_map(|context| {
 581                context
 582                    .iter()
 583                    .filter_map(|context_id| self.context.get(&context_id))
 584            })
 585    }
 586
 587    /// Returns whether all of the tool uses have finished running.
 588    pub fn all_tools_finished(&self) -> bool {
 589        // If the only pending tool uses left are the ones with errors, then
 590        // that means that we've finished running all of the pending tools.
 591        self.tool_use
 592            .pending_tool_uses()
 593            .iter()
 594            .all(|tool_use| tool_use.status.is_error())
 595    }
 596
 597    pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
 598        self.tool_use.tool_uses_for_message(id, cx)
 599    }
 600
 601    pub fn tool_results_for_message(&self, id: MessageId) -> Vec<&LanguageModelToolResult> {
 602        self.tool_use.tool_results_for_message(id)
 603    }
 604
 605    pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
 606        self.tool_use.tool_result(id)
 607    }
 608
 609    pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
 610        self.tool_use.message_has_tool_results(message_id)
 611    }
 612
 613    pub fn insert_user_message(
 614        &mut self,
 615        text: impl Into<String>,
 616        context: Vec<AssistantContext>,
 617        git_checkpoint: Option<GitStoreCheckpoint>,
 618        cx: &mut Context<Self>,
 619    ) -> MessageId {
 620        let text = text.into();
 621
 622        let message_id = self.insert_message(Role::User, vec![MessageSegment::Text(text)], cx);
 623
 624        // Filter out contexts that have already been included in previous messages
 625        let new_context: Vec<_> = context
 626            .into_iter()
 627            .filter(|ctx| !self.context.contains_key(&ctx.id()))
 628            .collect();
 629
 630        if !new_context.is_empty() {
 631            if let Some(context_string) = format_context_as_string(new_context.iter(), cx) {
 632                if let Some(message) = self.messages.iter_mut().find(|m| m.id == message_id) {
 633                    message.context = context_string;
 634                }
 635            }
 636
 637            self.action_log.update(cx, |log, cx| {
 638                // Track all buffers added as context
 639                for ctx in &new_context {
 640                    match ctx {
 641                        AssistantContext::File(file_ctx) => {
 642                            log.buffer_added_as_context(file_ctx.context_buffer.buffer.clone(), cx);
 643                        }
 644                        AssistantContext::Directory(dir_ctx) => {
 645                            for context_buffer in &dir_ctx.context_buffers {
 646                                log.buffer_added_as_context(context_buffer.buffer.clone(), cx);
 647                            }
 648                        }
 649                        AssistantContext::Symbol(symbol_ctx) => {
 650                            log.buffer_added_as_context(
 651                                symbol_ctx.context_symbol.buffer.clone(),
 652                                cx,
 653                            );
 654                        }
 655                        AssistantContext::FetchedUrl(_) | AssistantContext::Thread(_) => {}
 656                    }
 657                }
 658            });
 659        }
 660
 661        let context_ids = new_context
 662            .iter()
 663            .map(|context| context.id())
 664            .collect::<Vec<_>>();
 665        self.context.extend(
 666            new_context
 667                .into_iter()
 668                .map(|context| (context.id(), context)),
 669        );
 670        self.context_by_message.insert(message_id, context_ids);
 671
 672        if let Some(git_checkpoint) = git_checkpoint {
 673            self.pending_checkpoint = Some(ThreadCheckpoint {
 674                message_id,
 675                git_checkpoint,
 676            });
 677        }
 678        message_id
 679    }
 680
 681    pub fn insert_message(
 682        &mut self,
 683        role: Role,
 684        segments: Vec<MessageSegment>,
 685        cx: &mut Context<Self>,
 686    ) -> MessageId {
 687        let id = self.next_message_id.post_inc();
 688        self.messages.push(Message {
 689            id,
 690            role,
 691            segments,
 692            context: String::new(),
 693        });
 694        self.touch_updated_at();
 695        cx.emit(ThreadEvent::MessageAdded(id));
 696        id
 697    }
 698
 699    pub fn edit_message(
 700        &mut self,
 701        id: MessageId,
 702        new_role: Role,
 703        new_segments: Vec<MessageSegment>,
 704        cx: &mut Context<Self>,
 705    ) -> bool {
 706        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
 707            return false;
 708        };
 709        message.role = new_role;
 710        message.segments = new_segments;
 711        self.touch_updated_at();
 712        cx.emit(ThreadEvent::MessageEdited(id));
 713        true
 714    }
 715
 716    pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
 717        let Some(index) = self.messages.iter().position(|message| message.id == id) else {
 718            return false;
 719        };
 720        self.messages.remove(index);
 721        self.context_by_message.remove(&id);
 722        self.touch_updated_at();
 723        cx.emit(ThreadEvent::MessageDeleted(id));
 724        true
 725    }
 726
 727    /// Returns the representation of this [`Thread`] in a textual form.
 728    ///
 729    /// This is the representation we use when attaching a thread as context to another thread.
 730    pub fn text(&self) -> String {
 731        let mut text = String::new();
 732
 733        for message in &self.messages {
 734            text.push_str(match message.role {
 735                language_model::Role::User => "User:",
 736                language_model::Role::Assistant => "Assistant:",
 737                language_model::Role::System => "System:",
 738            });
 739            text.push('\n');
 740
 741            for segment in &message.segments {
 742                match segment {
 743                    MessageSegment::Text(content) => text.push_str(content),
 744                    MessageSegment::Thinking(content) => {
 745                        text.push_str(&format!("<think>{}</think>", content))
 746                    }
 747                }
 748            }
 749            text.push('\n');
 750        }
 751
 752        text
 753    }
 754
 755    /// Serializes this thread into a format for storage or telemetry.
 756    pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
 757        let initial_project_snapshot = self.initial_project_snapshot.clone();
 758        cx.spawn(async move |this, cx| {
 759            let initial_project_snapshot = initial_project_snapshot.await;
 760            this.read_with(cx, |this, cx| SerializedThread {
 761                version: SerializedThread::VERSION.to_string(),
 762                summary: this.summary_or_default(),
 763                updated_at: this.updated_at(),
 764                messages: this
 765                    .messages()
 766                    .map(|message| SerializedMessage {
 767                        id: message.id,
 768                        role: message.role,
 769                        segments: message
 770                            .segments
 771                            .iter()
 772                            .map(|segment| match segment {
 773                                MessageSegment::Text(text) => {
 774                                    SerializedMessageSegment::Text { text: text.clone() }
 775                                }
 776                                MessageSegment::Thinking(text) => {
 777                                    SerializedMessageSegment::Thinking { text: text.clone() }
 778                                }
 779                            })
 780                            .collect(),
 781                        tool_uses: this
 782                            .tool_uses_for_message(message.id, cx)
 783                            .into_iter()
 784                            .map(|tool_use| SerializedToolUse {
 785                                id: tool_use.id,
 786                                name: tool_use.name,
 787                                input: tool_use.input,
 788                            })
 789                            .collect(),
 790                        tool_results: this
 791                            .tool_results_for_message(message.id)
 792                            .into_iter()
 793                            .map(|tool_result| SerializedToolResult {
 794                                tool_use_id: tool_result.tool_use_id.clone(),
 795                                is_error: tool_result.is_error,
 796                                content: tool_result.content.clone(),
 797                            })
 798                            .collect(),
 799                        context: message.context.clone(),
 800                    })
 801                    .collect(),
 802                initial_project_snapshot,
 803                cumulative_token_usage: this.cumulative_token_usage.clone(),
 804                detailed_summary_state: this.detailed_summary_state.clone(),
 805            })
 806        })
 807    }
 808
 809    pub fn set_system_prompt_context(&mut self, context: AssistantSystemPromptContext) {
 810        self.system_prompt_context = Some(context);
 811    }
 812
 813    pub fn system_prompt_context(&self) -> &Option<AssistantSystemPromptContext> {
 814        &self.system_prompt_context
 815    }
 816
 817    pub fn load_system_prompt_context(
 818        &self,
 819        cx: &App,
 820    ) -> Task<(AssistantSystemPromptContext, Option<ThreadError>)> {
 821        let project = self.project.read(cx);
 822        let tasks = project
 823            .visible_worktrees(cx)
 824            .map(|worktree| {
 825                Self::load_worktree_info_for_system_prompt(
 826                    project.fs().clone(),
 827                    worktree.read(cx),
 828                    cx,
 829                )
 830            })
 831            .collect::<Vec<_>>();
 832
 833        cx.spawn(async |_cx| {
 834            let results = futures::future::join_all(tasks).await;
 835            let mut first_err = None;
 836            let worktrees = results
 837                .into_iter()
 838                .map(|(worktree, err)| {
 839                    if first_err.is_none() && err.is_some() {
 840                        first_err = err;
 841                    }
 842                    worktree
 843                })
 844                .collect::<Vec<_>>();
 845            (AssistantSystemPromptContext::new(worktrees), first_err)
 846        })
 847    }
 848
 849    fn load_worktree_info_for_system_prompt(
 850        fs: Arc<dyn Fs>,
 851        worktree: &Worktree,
 852        cx: &App,
 853    ) -> Task<(WorktreeInfoForSystemPrompt, Option<ThreadError>)> {
 854        let root_name = worktree.root_name().into();
 855        let abs_path = worktree.abs_path();
 856
 857        // Note that Cline supports `.clinerules` being a directory, but that is not currently
 858        // supported. This doesn't seem to occur often in GitHub repositories.
 859        const RULES_FILE_NAMES: [&'static str; 6] = [
 860            ".rules",
 861            ".cursorrules",
 862            ".windsurfrules",
 863            ".clinerules",
 864            ".github/copilot-instructions.md",
 865            "CLAUDE.md",
 866        ];
 867        let selected_rules_file = RULES_FILE_NAMES
 868            .into_iter()
 869            .filter_map(|name| {
 870                worktree
 871                    .entry_for_path(name)
 872                    .filter(|entry| entry.is_file())
 873                    .map(|entry| (entry.path.clone(), worktree.absolutize(&entry.path)))
 874            })
 875            .next();
 876
 877        if let Some((rel_rules_path, abs_rules_path)) = selected_rules_file {
 878            cx.spawn(async move |_| {
 879                let rules_file_result = maybe!(async move {
 880                    let abs_rules_path = abs_rules_path?;
 881                    let text = fs.load(&abs_rules_path).await.with_context(|| {
 882                        format!("Failed to load assistant rules file {:?}", abs_rules_path)
 883                    })?;
 884                    anyhow::Ok(RulesFile {
 885                        rel_path: rel_rules_path,
 886                        abs_path: abs_rules_path.into(),
 887                        text: text.trim().to_string(),
 888                    })
 889                })
 890                .await;
 891                let (rules_file, rules_file_error) = match rules_file_result {
 892                    Ok(rules_file) => (Some(rules_file), None),
 893                    Err(err) => (
 894                        None,
 895                        Some(ThreadError::Message {
 896                            header: "Error loading rules file".into(),
 897                            message: format!("{err}").into(),
 898                        }),
 899                    ),
 900                };
 901                let worktree_info = WorktreeInfoForSystemPrompt {
 902                    root_name,
 903                    abs_path,
 904                    rules_file,
 905                };
 906                (worktree_info, rules_file_error)
 907            })
 908        } else {
 909            Task::ready((
 910                WorktreeInfoForSystemPrompt {
 911                    root_name,
 912                    abs_path,
 913                    rules_file: None,
 914                },
 915                None,
 916            ))
 917        }
 918    }
 919
 920    pub fn send_to_model(
 921        &mut self,
 922        model: Arc<dyn LanguageModel>,
 923        request_kind: RequestKind,
 924        cx: &mut Context<Self>,
 925    ) {
 926        let mut request = self.to_completion_request(request_kind, cx);
 927        if model.supports_tools() {
 928            request.tools = {
 929                let mut tools = Vec::new();
 930                tools.extend(self.tools().enabled_tools(cx).into_iter().map(|tool| {
 931                    LanguageModelRequestTool {
 932                        name: tool.name(),
 933                        description: tool.description(),
 934                        input_schema: tool.input_schema(model.tool_input_format()),
 935                    }
 936                }));
 937
 938                tools
 939            };
 940        }
 941
 942        self.stream_completion(request, model, cx);
 943    }
 944
 945    pub fn used_tools_since_last_user_message(&self) -> bool {
 946        for message in self.messages.iter().rev() {
 947            if self.tool_use.message_has_tool_results(message.id) {
 948                return true;
 949            } else if message.role == Role::User {
 950                return false;
 951            }
 952        }
 953
 954        false
 955    }
 956
 957    pub fn to_completion_request(
 958        &self,
 959        request_kind: RequestKind,
 960        cx: &App,
 961    ) -> LanguageModelRequest {
 962        let mut request = LanguageModelRequest {
 963            messages: vec![],
 964            tools: Vec::new(),
 965            stop: Vec::new(),
 966            temperature: None,
 967        };
 968
 969        if let Some(system_prompt_context) = self.system_prompt_context.as_ref() {
 970            if let Some(system_prompt) = self
 971                .prompt_builder
 972                .generate_assistant_system_prompt(system_prompt_context)
 973                .context("failed to generate assistant system prompt")
 974                .log_err()
 975            {
 976                request.messages.push(LanguageModelRequestMessage {
 977                    role: Role::System,
 978                    content: vec![MessageContent::Text(system_prompt)],
 979                    cache: true,
 980                });
 981            }
 982        } else {
 983            log::error!("system_prompt_context not set.")
 984        }
 985
 986        for message in &self.messages {
 987            let mut request_message = LanguageModelRequestMessage {
 988                role: message.role,
 989                content: Vec::new(),
 990                cache: false,
 991            };
 992
 993            match request_kind {
 994                RequestKind::Chat => {
 995                    self.tool_use
 996                        .attach_tool_results(message.id, &mut request_message);
 997                }
 998                RequestKind::Summarize => {
 999                    // We don't care about tool use during summarization.
1000                    if self.tool_use.message_has_tool_results(message.id) {
1001                        continue;
1002                    }
1003                }
1004            }
1005
1006            if !message.segments.is_empty() {
1007                request_message
1008                    .content
1009                    .push(MessageContent::Text(message.to_string()));
1010            }
1011
1012            match request_kind {
1013                RequestKind::Chat => {
1014                    self.tool_use
1015                        .attach_tool_uses(message.id, &mut request_message);
1016                }
1017                RequestKind::Summarize => {
1018                    // We don't care about tool use during summarization.
1019                }
1020            };
1021
1022            request.messages.push(request_message);
1023        }
1024
1025        // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
1026        if let Some(last) = request.messages.last_mut() {
1027            last.cache = true;
1028        }
1029
1030        self.attached_tracked_files_state(&mut request.messages, cx);
1031
1032        request
1033    }
1034
1035    fn attached_tracked_files_state(
1036        &self,
1037        messages: &mut Vec<LanguageModelRequestMessage>,
1038        cx: &App,
1039    ) {
1040        const STALE_FILES_HEADER: &str = "These files changed since last read:";
1041
1042        let mut stale_message = String::new();
1043
1044        let action_log = self.action_log.read(cx);
1045
1046        for stale_file in action_log.stale_buffers(cx) {
1047            let Some(file) = stale_file.read(cx).file() else {
1048                continue;
1049            };
1050
1051            if stale_message.is_empty() {
1052                write!(&mut stale_message, "{}\n", STALE_FILES_HEADER).ok();
1053            }
1054
1055            writeln!(&mut stale_message, "- {}", file.path().display()).ok();
1056        }
1057
1058        let mut content = Vec::with_capacity(2);
1059
1060        if !stale_message.is_empty() {
1061            content.push(stale_message.into());
1062        }
1063
1064        if action_log.has_edited_files_since_project_diagnostics_check() {
1065            content.push(
1066                "\n\nWhen you're done making changes, make sure to check project diagnostics \
1067                and fix all errors AND warnings you introduced! \
1068                DO NOT mention you're going to do this until you're done."
1069                    .into(),
1070            );
1071        }
1072
1073        if !content.is_empty() {
1074            let context_message = LanguageModelRequestMessage {
1075                role: Role::User,
1076                content,
1077                cache: false,
1078            };
1079
1080            messages.push(context_message);
1081        }
1082    }
1083
1084    pub fn stream_completion(
1085        &mut self,
1086        request: LanguageModelRequest,
1087        model: Arc<dyn LanguageModel>,
1088        cx: &mut Context<Self>,
1089    ) {
1090        let pending_completion_id = post_inc(&mut self.completion_count);
1091
1092        let task = cx.spawn(async move |thread, cx| {
1093            let stream = model.stream_completion(request, &cx);
1094            let initial_token_usage =
1095                thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage.clone());
1096            let stream_completion = async {
1097                let mut events = stream.await?;
1098                let mut stop_reason = StopReason::EndTurn;
1099                let mut current_token_usage = TokenUsage::default();
1100
1101                while let Some(event) = events.next().await {
1102                    let event = event?;
1103
1104                    thread.update(cx, |thread, cx| {
1105                        match event {
1106                            LanguageModelCompletionEvent::StartMessage { .. } => {
1107                                thread.insert_message(
1108                                    Role::Assistant,
1109                                    vec![MessageSegment::Text(String::new())],
1110                                    cx,
1111                                );
1112                            }
1113                            LanguageModelCompletionEvent::Stop(reason) => {
1114                                stop_reason = reason;
1115                            }
1116                            LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1117                                thread.cumulative_token_usage =
1118                                    thread.cumulative_token_usage.clone() + token_usage.clone()
1119                                        - current_token_usage.clone();
1120                                current_token_usage = token_usage;
1121                            }
1122                            LanguageModelCompletionEvent::Text(chunk) => {
1123                                if let Some(last_message) = thread.messages.last_mut() {
1124                                    if last_message.role == Role::Assistant {
1125                                        last_message.push_text(&chunk);
1126                                        cx.emit(ThreadEvent::StreamedAssistantText(
1127                                            last_message.id,
1128                                            chunk,
1129                                        ));
1130                                    } else {
1131                                        // If we won't have an Assistant message yet, assume this chunk marks the beginning
1132                                        // of a new Assistant response.
1133                                        //
1134                                        // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1135                                        // will result in duplicating the text of the chunk in the rendered Markdown.
1136                                        thread.insert_message(
1137                                            Role::Assistant,
1138                                            vec![MessageSegment::Text(chunk.to_string())],
1139                                            cx,
1140                                        );
1141                                    };
1142                                }
1143                            }
1144                            LanguageModelCompletionEvent::Thinking(chunk) => {
1145                                if let Some(last_message) = thread.messages.last_mut() {
1146                                    if last_message.role == Role::Assistant {
1147                                        last_message.push_thinking(&chunk);
1148                                        cx.emit(ThreadEvent::StreamedAssistantThinking(
1149                                            last_message.id,
1150                                            chunk,
1151                                        ));
1152                                    } else {
1153                                        // If we won't have an Assistant message yet, assume this chunk marks the beginning
1154                                        // of a new Assistant response.
1155                                        //
1156                                        // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1157                                        // will result in duplicating the text of the chunk in the rendered Markdown.
1158                                        thread.insert_message(
1159                                            Role::Assistant,
1160                                            vec![MessageSegment::Thinking(chunk.to_string())],
1161                                            cx,
1162                                        );
1163                                    };
1164                                }
1165                            }
1166                            LanguageModelCompletionEvent::ToolUse(tool_use) => {
1167                                let last_assistant_message_id = thread
1168                                    .messages
1169                                    .iter_mut()
1170                                    .rfind(|message| message.role == Role::Assistant)
1171                                    .map(|message| message.id)
1172                                    .unwrap_or_else(|| {
1173                                        thread.insert_message(Role::Assistant, vec![], cx)
1174                                    });
1175
1176                                thread.tool_use.request_tool_use(
1177                                    last_assistant_message_id,
1178                                    tool_use,
1179                                    cx,
1180                                );
1181                            }
1182                        }
1183
1184                        thread.touch_updated_at();
1185                        cx.emit(ThreadEvent::StreamedCompletion);
1186                        cx.notify();
1187                    })?;
1188
1189                    smol::future::yield_now().await;
1190                }
1191
1192                thread.update(cx, |thread, cx| {
1193                    thread
1194                        .pending_completions
1195                        .retain(|completion| completion.id != pending_completion_id);
1196
1197                    if thread.summary.is_none() && thread.messages.len() >= 2 {
1198                        thread.summarize(cx);
1199                    }
1200                })?;
1201
1202                anyhow::Ok(stop_reason)
1203            };
1204
1205            let result = stream_completion.await;
1206
1207            thread
1208                .update(cx, |thread, cx| {
1209                    thread.finalize_pending_checkpoint(cx);
1210                    match result.as_ref() {
1211                        Ok(stop_reason) => match stop_reason {
1212                            StopReason::ToolUse => {
1213                                cx.emit(ThreadEvent::UsePendingTools);
1214                            }
1215                            StopReason::EndTurn => {}
1216                            StopReason::MaxTokens => {}
1217                        },
1218                        Err(error) => {
1219                            if error.is::<PaymentRequiredError>() {
1220                                cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
1221                            } else if error.is::<MaxMonthlySpendReachedError>() {
1222                                cx.emit(ThreadEvent::ShowError(
1223                                    ThreadError::MaxMonthlySpendReached,
1224                                ));
1225                            } else {
1226                                let error_message = error
1227                                    .chain()
1228                                    .map(|err| err.to_string())
1229                                    .collect::<Vec<_>>()
1230                                    .join("\n");
1231                                cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1232                                    header: "Error interacting with language model".into(),
1233                                    message: SharedString::from(error_message.clone()),
1234                                }));
1235                            }
1236
1237                            thread.cancel_last_completion(cx);
1238                        }
1239                    }
1240                    cx.emit(ThreadEvent::DoneStreaming);
1241
1242                    if let Ok(initial_usage) = initial_token_usage {
1243                        let usage = thread.cumulative_token_usage.clone() - initial_usage;
1244
1245                        telemetry::event!(
1246                            "Assistant Thread Completion",
1247                            thread_id = thread.id().to_string(),
1248                            model = model.telemetry_id(),
1249                            model_provider = model.provider_id().to_string(),
1250                            input_tokens = usage.input_tokens,
1251                            output_tokens = usage.output_tokens,
1252                            cache_creation_input_tokens = usage.cache_creation_input_tokens,
1253                            cache_read_input_tokens = usage.cache_read_input_tokens,
1254                        );
1255                    }
1256                })
1257                .ok();
1258        });
1259
1260        self.pending_completions.push(PendingCompletion {
1261            id: pending_completion_id,
1262            _task: task,
1263        });
1264    }
1265
1266    pub fn summarize(&mut self, cx: &mut Context<Self>) {
1267        let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
1268            return;
1269        };
1270
1271        if !model.provider.is_authenticated(cx) {
1272            return;
1273        }
1274
1275        let mut request = self.to_completion_request(RequestKind::Summarize, cx);
1276        request.messages.push(LanguageModelRequestMessage {
1277            role: Role::User,
1278            content: vec![
1279                "Generate a concise 3-7 word title for this conversation, omitting punctuation. \
1280                 Go straight to the title, without any preamble and prefix like `Here's a concise suggestion:...` or `Title:`. \
1281                 If the conversation is about a specific subject, include it in the title. \
1282                 Be descriptive. DO NOT speak in the first person."
1283                    .into(),
1284            ],
1285            cache: false,
1286        });
1287
1288        self.pending_summary = cx.spawn(async move |this, cx| {
1289            async move {
1290                let stream = model.model.stream_completion_text(request, &cx);
1291                let mut messages = stream.await?;
1292
1293                let mut new_summary = String::new();
1294                while let Some(message) = messages.stream.next().await {
1295                    let text = message?;
1296                    let mut lines = text.lines();
1297                    new_summary.extend(lines.next());
1298
1299                    // Stop if the LLM generated multiple lines.
1300                    if lines.next().is_some() {
1301                        break;
1302                    }
1303                }
1304
1305                this.update(cx, |this, cx| {
1306                    if !new_summary.is_empty() {
1307                        this.summary = Some(new_summary.into());
1308                    }
1309
1310                    cx.emit(ThreadEvent::SummaryGenerated);
1311                })?;
1312
1313                anyhow::Ok(())
1314            }
1315            .log_err()
1316            .await
1317        });
1318    }
1319
1320    pub fn generate_detailed_summary(&mut self, cx: &mut Context<Self>) -> Option<Task<()>> {
1321        let last_message_id = self.messages.last().map(|message| message.id)?;
1322
1323        match &self.detailed_summary_state {
1324            DetailedSummaryState::Generating { message_id, .. }
1325            | DetailedSummaryState::Generated { message_id, .. }
1326                if *message_id == last_message_id =>
1327            {
1328                // Already up-to-date
1329                return None;
1330            }
1331            _ => {}
1332        }
1333
1334        let ConfiguredModel { model, provider } =
1335            LanguageModelRegistry::read_global(cx).thread_summary_model()?;
1336
1337        if !provider.is_authenticated(cx) {
1338            return None;
1339        }
1340
1341        let mut request = self.to_completion_request(RequestKind::Summarize, cx);
1342
1343        request.messages.push(LanguageModelRequestMessage {
1344            role: Role::User,
1345            content: vec![
1346                "Generate a detailed summary of this conversation. Include:\n\
1347                1. A brief overview of what was discussed\n\
1348                2. Key facts or information discovered\n\
1349                3. Outcomes or conclusions reached\n\
1350                4. Any action items or next steps if any\n\
1351                Format it in Markdown with headings and bullet points."
1352                    .into(),
1353            ],
1354            cache: false,
1355        });
1356
1357        let task = cx.spawn(async move |thread, cx| {
1358            let stream = model.stream_completion_text(request, &cx);
1359            let Some(mut messages) = stream.await.log_err() else {
1360                thread
1361                    .update(cx, |this, _cx| {
1362                        this.detailed_summary_state = DetailedSummaryState::NotGenerated;
1363                    })
1364                    .log_err();
1365
1366                return;
1367            };
1368
1369            let mut new_detailed_summary = String::new();
1370
1371            while let Some(chunk) = messages.stream.next().await {
1372                if let Some(chunk) = chunk.log_err() {
1373                    new_detailed_summary.push_str(&chunk);
1374                }
1375            }
1376
1377            thread
1378                .update(cx, |this, _cx| {
1379                    this.detailed_summary_state = DetailedSummaryState::Generated {
1380                        text: new_detailed_summary.into(),
1381                        message_id: last_message_id,
1382                    };
1383                })
1384                .log_err();
1385        });
1386
1387        self.detailed_summary_state = DetailedSummaryState::Generating {
1388            message_id: last_message_id,
1389        };
1390
1391        Some(task)
1392    }
1393
1394    pub fn is_generating_detailed_summary(&self) -> bool {
1395        matches!(
1396            self.detailed_summary_state,
1397            DetailedSummaryState::Generating { .. }
1398        )
1399    }
1400
1401    pub fn use_pending_tools(
1402        &mut self,
1403        cx: &mut Context<Self>,
1404    ) -> impl IntoIterator<Item = PendingToolUse> + use<> {
1405        let request = self.to_completion_request(RequestKind::Chat, cx);
1406        let messages = Arc::new(request.messages);
1407        let pending_tool_uses = self
1408            .tool_use
1409            .pending_tool_uses()
1410            .into_iter()
1411            .filter(|tool_use| tool_use.status.is_idle())
1412            .cloned()
1413            .collect::<Vec<_>>();
1414
1415        for tool_use in pending_tool_uses.iter() {
1416            if let Some(tool) = self.tools.tool(&tool_use.name, cx) {
1417                if tool.needs_confirmation(&tool_use.input, cx)
1418                    && !AssistantSettings::get_global(cx).always_allow_tool_actions
1419                {
1420                    self.tool_use.confirm_tool_use(
1421                        tool_use.id.clone(),
1422                        tool_use.ui_text.clone(),
1423                        tool_use.input.clone(),
1424                        messages.clone(),
1425                        tool,
1426                    );
1427                    cx.emit(ThreadEvent::ToolConfirmationNeeded);
1428                } else {
1429                    self.run_tool(
1430                        tool_use.id.clone(),
1431                        tool_use.ui_text.clone(),
1432                        tool_use.input.clone(),
1433                        &messages,
1434                        tool,
1435                        cx,
1436                    );
1437                }
1438            }
1439        }
1440
1441        pending_tool_uses
1442    }
1443
1444    pub fn run_tool(
1445        &mut self,
1446        tool_use_id: LanguageModelToolUseId,
1447        ui_text: impl Into<SharedString>,
1448        input: serde_json::Value,
1449        messages: &[LanguageModelRequestMessage],
1450        tool: Arc<dyn Tool>,
1451        cx: &mut Context<Thread>,
1452    ) {
1453        let task = self.spawn_tool_use(tool_use_id.clone(), messages, input, tool, cx);
1454        self.tool_use
1455            .run_pending_tool(tool_use_id, ui_text.into(), task);
1456    }
1457
1458    fn spawn_tool_use(
1459        &mut self,
1460        tool_use_id: LanguageModelToolUseId,
1461        messages: &[LanguageModelRequestMessage],
1462        input: serde_json::Value,
1463        tool: Arc<dyn Tool>,
1464        cx: &mut Context<Thread>,
1465    ) -> Task<()> {
1466        let tool_name: Arc<str> = tool.name().into();
1467
1468        let run_tool = if self.tools.is_disabled(&tool.source(), &tool_name) {
1469            Task::ready(Err(anyhow!("tool is disabled: {tool_name}")))
1470        } else {
1471            tool.run(
1472                input,
1473                messages,
1474                self.project.clone(),
1475                self.action_log.clone(),
1476                cx,
1477            )
1478        };
1479
1480        cx.spawn({
1481            async move |thread: WeakEntity<Thread>, cx| {
1482                let output = run_tool.await;
1483
1484                thread
1485                    .update(cx, |thread, cx| {
1486                        let pending_tool_use = thread.tool_use.insert_tool_output(
1487                            tool_use_id.clone(),
1488                            tool_name,
1489                            output,
1490                            cx,
1491                        );
1492
1493                        cx.emit(ThreadEvent::ToolFinished {
1494                            tool_use_id,
1495                            pending_tool_use,
1496                            canceled: false,
1497                        });
1498                    })
1499                    .ok();
1500            }
1501        })
1502    }
1503
1504    pub fn attach_tool_results(&mut self, cx: &mut Context<Self>) {
1505        // Insert a user message to contain the tool results.
1506        self.insert_user_message(
1507            // TODO: Sending up a user message without any content results in the model sending back
1508            // responses that also don't have any content. We currently don't handle this case well,
1509            // so for now we provide some text to keep the model on track.
1510            "Here are the tool results.",
1511            Vec::new(),
1512            None,
1513            cx,
1514        );
1515    }
1516
1517    /// Cancels the last pending completion, if there are any pending.
1518    ///
1519    /// Returns whether a completion was canceled.
1520    pub fn cancel_last_completion(&mut self, cx: &mut Context<Self>) -> bool {
1521        let canceled = if self.pending_completions.pop().is_some() {
1522            true
1523        } else {
1524            let mut canceled = false;
1525            for pending_tool_use in self.tool_use.cancel_pending() {
1526                canceled = true;
1527                cx.emit(ThreadEvent::ToolFinished {
1528                    tool_use_id: pending_tool_use.id.clone(),
1529                    pending_tool_use: Some(pending_tool_use),
1530                    canceled: true,
1531                });
1532            }
1533            canceled
1534        };
1535        self.finalize_pending_checkpoint(cx);
1536        canceled
1537    }
1538
1539    /// Returns the feedback given to the thread, if any.
1540    pub fn feedback(&self) -> Option<ThreadFeedback> {
1541        self.feedback
1542    }
1543
1544    /// Reports feedback about the thread and stores it in our telemetry backend.
1545    pub fn report_feedback(
1546        &mut self,
1547        feedback: ThreadFeedback,
1548        cx: &mut Context<Self>,
1549    ) -> Task<Result<()>> {
1550        let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
1551        let serialized_thread = self.serialize(cx);
1552        let thread_id = self.id().clone();
1553        let client = self.project.read(cx).client();
1554        self.feedback = Some(feedback);
1555        cx.notify();
1556
1557        cx.background_spawn(async move {
1558            let final_project_snapshot = final_project_snapshot.await;
1559            let serialized_thread = serialized_thread.await?;
1560            let thread_data =
1561                serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
1562
1563            let rating = match feedback {
1564                ThreadFeedback::Positive => "positive",
1565                ThreadFeedback::Negative => "negative",
1566            };
1567            telemetry::event!(
1568                "Assistant Thread Rated",
1569                rating,
1570                thread_id,
1571                thread_data,
1572                final_project_snapshot
1573            );
1574            client.telemetry().flush_events();
1575
1576            Ok(())
1577        })
1578    }
1579
1580    /// Create a snapshot of the current project state including git information and unsaved buffers.
1581    fn project_snapshot(
1582        project: Entity<Project>,
1583        cx: &mut Context<Self>,
1584    ) -> Task<Arc<ProjectSnapshot>> {
1585        let git_store = project.read(cx).git_store().clone();
1586        let worktree_snapshots: Vec<_> = project
1587            .read(cx)
1588            .visible_worktrees(cx)
1589            .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
1590            .collect();
1591
1592        cx.spawn(async move |_, cx| {
1593            let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
1594
1595            let mut unsaved_buffers = Vec::new();
1596            cx.update(|app_cx| {
1597                let buffer_store = project.read(app_cx).buffer_store();
1598                for buffer_handle in buffer_store.read(app_cx).buffers() {
1599                    let buffer = buffer_handle.read(app_cx);
1600                    if buffer.is_dirty() {
1601                        if let Some(file) = buffer.file() {
1602                            let path = file.path().to_string_lossy().to_string();
1603                            unsaved_buffers.push(path);
1604                        }
1605                    }
1606                }
1607            })
1608            .ok();
1609
1610            Arc::new(ProjectSnapshot {
1611                worktree_snapshots,
1612                unsaved_buffer_paths: unsaved_buffers,
1613                timestamp: Utc::now(),
1614            })
1615        })
1616    }
1617
1618    fn worktree_snapshot(
1619        worktree: Entity<project::Worktree>,
1620        git_store: Entity<GitStore>,
1621        cx: &App,
1622    ) -> Task<WorktreeSnapshot> {
1623        cx.spawn(async move |cx| {
1624            // Get worktree path and snapshot
1625            let worktree_info = cx.update(|app_cx| {
1626                let worktree = worktree.read(app_cx);
1627                let path = worktree.abs_path().to_string_lossy().to_string();
1628                let snapshot = worktree.snapshot();
1629                (path, snapshot)
1630            });
1631
1632            let Ok((worktree_path, _snapshot)) = worktree_info else {
1633                return WorktreeSnapshot {
1634                    worktree_path: String::new(),
1635                    git_state: None,
1636                };
1637            };
1638
1639            let git_state = git_store
1640                .update(cx, |git_store, cx| {
1641                    git_store
1642                        .repositories()
1643                        .values()
1644                        .find(|repo| {
1645                            repo.read(cx)
1646                                .abs_path_to_repo_path(&worktree.read(cx).abs_path())
1647                                .is_some()
1648                        })
1649                        .cloned()
1650                })
1651                .ok()
1652                .flatten()
1653                .map(|repo| {
1654                    repo.update(cx, |repo, _| {
1655                        let current_branch =
1656                            repo.branch.as_ref().map(|branch| branch.name.to_string());
1657                        repo.send_job(None, |state, _| async move {
1658                            let RepositoryState::Local { backend, .. } = state else {
1659                                return GitState {
1660                                    remote_url: None,
1661                                    head_sha: None,
1662                                    current_branch,
1663                                    diff: None,
1664                                };
1665                            };
1666
1667                            let remote_url = backend.remote_url("origin");
1668                            let head_sha = backend.head_sha();
1669                            let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
1670
1671                            GitState {
1672                                remote_url,
1673                                head_sha,
1674                                current_branch,
1675                                diff,
1676                            }
1677                        })
1678                    })
1679                });
1680
1681            let git_state = match git_state {
1682                Some(git_state) => match git_state.ok() {
1683                    Some(git_state) => git_state.await.ok(),
1684                    None => None,
1685                },
1686                None => None,
1687            };
1688
1689            WorktreeSnapshot {
1690                worktree_path,
1691                git_state,
1692            }
1693        })
1694    }
1695
1696    pub fn to_markdown(&self, cx: &App) -> Result<String> {
1697        let mut markdown = Vec::new();
1698
1699        if let Some(summary) = self.summary() {
1700            writeln!(markdown, "# {summary}\n")?;
1701        };
1702
1703        for message in self.messages() {
1704            writeln!(
1705                markdown,
1706                "## {role}\n",
1707                role = match message.role {
1708                    Role::User => "User",
1709                    Role::Assistant => "Assistant",
1710                    Role::System => "System",
1711                }
1712            )?;
1713
1714            if !message.context.is_empty() {
1715                writeln!(markdown, "{}", message.context)?;
1716            }
1717
1718            for segment in &message.segments {
1719                match segment {
1720                    MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
1721                    MessageSegment::Thinking(text) => {
1722                        writeln!(markdown, "<think>{}</think>\n", text)?
1723                    }
1724                }
1725            }
1726
1727            for tool_use in self.tool_uses_for_message(message.id, cx) {
1728                writeln!(
1729                    markdown,
1730                    "**Use Tool: {} ({})**",
1731                    tool_use.name, tool_use.id
1732                )?;
1733                writeln!(markdown, "```json")?;
1734                writeln!(
1735                    markdown,
1736                    "{}",
1737                    serde_json::to_string_pretty(&tool_use.input)?
1738                )?;
1739                writeln!(markdown, "```")?;
1740            }
1741
1742            for tool_result in self.tool_results_for_message(message.id) {
1743                write!(markdown, "**Tool Results: {}", tool_result.tool_use_id)?;
1744                if tool_result.is_error {
1745                    write!(markdown, " (Error)")?;
1746                }
1747
1748                writeln!(markdown, "**\n")?;
1749                writeln!(markdown, "{}", tool_result.content)?;
1750            }
1751        }
1752
1753        Ok(String::from_utf8_lossy(&markdown).to_string())
1754    }
1755
1756    pub fn keep_edits_in_range(
1757        &mut self,
1758        buffer: Entity<language::Buffer>,
1759        buffer_range: Range<language::Anchor>,
1760        cx: &mut Context<Self>,
1761    ) {
1762        self.action_log.update(cx, |action_log, cx| {
1763            action_log.keep_edits_in_range(buffer, buffer_range, cx)
1764        });
1765    }
1766
1767    pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
1768        self.action_log
1769            .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
1770    }
1771
1772    pub fn reject_edits_in_range(
1773        &mut self,
1774        buffer: Entity<language::Buffer>,
1775        buffer_range: Range<language::Anchor>,
1776        cx: &mut Context<Self>,
1777    ) -> Task<Result<()>> {
1778        self.action_log.update(cx, |action_log, cx| {
1779            action_log.reject_edits_in_range(buffer, buffer_range, cx)
1780        })
1781    }
1782
1783    pub fn action_log(&self) -> &Entity<ActionLog> {
1784        &self.action_log
1785    }
1786
1787    pub fn project(&self) -> &Entity<Project> {
1788        &self.project
1789    }
1790
1791    pub fn cumulative_token_usage(&self) -> TokenUsage {
1792        self.cumulative_token_usage.clone()
1793    }
1794
1795    pub fn total_token_usage(&self, cx: &App) -> TotalTokenUsage {
1796        let model_registry = LanguageModelRegistry::read_global(cx);
1797        let Some(model) = model_registry.default_model() else {
1798            return TotalTokenUsage::default();
1799        };
1800
1801        let max = model.model.max_token_count();
1802
1803        #[cfg(debug_assertions)]
1804        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
1805            .unwrap_or("0.8".to_string())
1806            .parse()
1807            .unwrap();
1808        #[cfg(not(debug_assertions))]
1809        let warning_threshold: f32 = 0.8;
1810
1811        let total = self.cumulative_token_usage.total_tokens() as usize;
1812
1813        let ratio = if total >= max {
1814            TokenUsageRatio::Exceeded
1815        } else if total as f32 / max as f32 >= warning_threshold {
1816            TokenUsageRatio::Warning
1817        } else {
1818            TokenUsageRatio::Normal
1819        };
1820
1821        TotalTokenUsage { total, max, ratio }
1822    }
1823
1824    pub fn deny_tool_use(
1825        &mut self,
1826        tool_use_id: LanguageModelToolUseId,
1827        tool_name: Arc<str>,
1828        cx: &mut Context<Self>,
1829    ) {
1830        let err = Err(anyhow::anyhow!(
1831            "Permission to run tool action denied by user"
1832        ));
1833
1834        self.tool_use
1835            .insert_tool_output(tool_use_id.clone(), tool_name, err, cx);
1836
1837        cx.emit(ThreadEvent::ToolFinished {
1838            tool_use_id,
1839            pending_tool_use: None,
1840            canceled: true,
1841        });
1842    }
1843}
1844
1845#[derive(Debug, Clone)]
1846pub enum ThreadError {
1847    PaymentRequired,
1848    MaxMonthlySpendReached,
1849    Message {
1850        header: SharedString,
1851        message: SharedString,
1852    },
1853}
1854
1855#[derive(Debug, Clone)]
1856pub enum ThreadEvent {
1857    ShowError(ThreadError),
1858    StreamedCompletion,
1859    StreamedAssistantText(MessageId, String),
1860    StreamedAssistantThinking(MessageId, String),
1861    DoneStreaming,
1862    MessageAdded(MessageId),
1863    MessageEdited(MessageId),
1864    MessageDeleted(MessageId),
1865    SummaryGenerated,
1866    SummaryChanged,
1867    UsePendingTools,
1868    ToolFinished {
1869        #[allow(unused)]
1870        tool_use_id: LanguageModelToolUseId,
1871        /// The pending tool use that corresponds to this tool.
1872        pending_tool_use: Option<PendingToolUse>,
1873        /// Whether the tool was canceled by the user.
1874        canceled: bool,
1875    },
1876    CheckpointChanged,
1877    ToolConfirmationNeeded,
1878}
1879
1880impl EventEmitter<ThreadEvent> for Thread {}
1881
1882struct PendingCompletion {
1883    id: usize,
1884    _task: Task<()>,
1885}
1886
1887#[cfg(test)]
1888mod tests {
1889    use super::*;
1890    use crate::{ThreadStore, context_store::ContextStore, thread_store};
1891    use assistant_settings::AssistantSettings;
1892    use context_server::ContextServerSettings;
1893    use editor::EditorSettings;
1894    use gpui::TestAppContext;
1895    use project::{FakeFs, Project};
1896    use prompt_store::PromptBuilder;
1897    use serde_json::json;
1898    use settings::{Settings, SettingsStore};
1899    use std::sync::Arc;
1900    use theme::ThemeSettings;
1901    use util::path;
1902    use workspace::Workspace;
1903
1904    #[gpui::test]
1905    async fn test_message_with_context(cx: &mut TestAppContext) {
1906        init_test_settings(cx);
1907
1908        let project = create_test_project(
1909            cx,
1910            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
1911        )
1912        .await;
1913
1914        let (_workspace, _thread_store, thread, context_store) =
1915            setup_test_environment(cx, project.clone()).await;
1916
1917        add_file_to_context(&project, &context_store, "test/code.rs", cx)
1918            .await
1919            .unwrap();
1920
1921        let context =
1922            context_store.update(cx, |store, _| store.context().first().cloned().unwrap());
1923
1924        // Insert user message with context
1925        let message_id = thread.update(cx, |thread, cx| {
1926            thread.insert_user_message("Please explain this code", vec![context], None, cx)
1927        });
1928
1929        // Check content and context in message object
1930        let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
1931
1932        // Use different path format strings based on platform for the test
1933        #[cfg(windows)]
1934        let path_part = r"test\code.rs";
1935        #[cfg(not(windows))]
1936        let path_part = "test/code.rs";
1937
1938        let expected_context = format!(
1939            r#"
1940<context>
1941The following items were attached by the user. You don't need to use other tools to read them.
1942
1943<files>
1944```rs {path_part}
1945fn main() {{
1946    println!("Hello, world!");
1947}}
1948```
1949</files>
1950</context>
1951"#
1952        );
1953
1954        assert_eq!(message.role, Role::User);
1955        assert_eq!(message.segments.len(), 1);
1956        assert_eq!(
1957            message.segments[0],
1958            MessageSegment::Text("Please explain this code".to_string())
1959        );
1960        assert_eq!(message.context, expected_context);
1961
1962        // Check message in request
1963        let request = thread.read_with(cx, |thread, cx| {
1964            thread.to_completion_request(RequestKind::Chat, cx)
1965        });
1966
1967        assert_eq!(request.messages.len(), 1);
1968        let expected_full_message = format!("{}Please explain this code", expected_context);
1969        assert_eq!(request.messages[0].string_contents(), expected_full_message);
1970    }
1971
1972    #[gpui::test]
1973    async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
1974        init_test_settings(cx);
1975
1976        let project = create_test_project(
1977            cx,
1978            json!({
1979                "file1.rs": "fn function1() {}\n",
1980                "file2.rs": "fn function2() {}\n",
1981                "file3.rs": "fn function3() {}\n",
1982            }),
1983        )
1984        .await;
1985
1986        let (_, _thread_store, thread, context_store) =
1987            setup_test_environment(cx, project.clone()).await;
1988
1989        // Open files individually
1990        add_file_to_context(&project, &context_store, "test/file1.rs", cx)
1991            .await
1992            .unwrap();
1993        add_file_to_context(&project, &context_store, "test/file2.rs", cx)
1994            .await
1995            .unwrap();
1996        add_file_to_context(&project, &context_store, "test/file3.rs", cx)
1997            .await
1998            .unwrap();
1999
2000        // Get the context objects
2001        let contexts = context_store.update(cx, |store, _| store.context().clone());
2002        assert_eq!(contexts.len(), 3);
2003
2004        // First message with context 1
2005        let message1_id = thread.update(cx, |thread, cx| {
2006            thread.insert_user_message("Message 1", vec![contexts[0].clone()], None, cx)
2007        });
2008
2009        // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
2010        let message2_id = thread.update(cx, |thread, cx| {
2011            thread.insert_user_message(
2012                "Message 2",
2013                vec![contexts[0].clone(), contexts[1].clone()],
2014                None,
2015                cx,
2016            )
2017        });
2018
2019        // Third message with all three contexts (contexts 1 and 2 should be skipped)
2020        let message3_id = thread.update(cx, |thread, cx| {
2021            thread.insert_user_message(
2022                "Message 3",
2023                vec![
2024                    contexts[0].clone(),
2025                    contexts[1].clone(),
2026                    contexts[2].clone(),
2027                ],
2028                None,
2029                cx,
2030            )
2031        });
2032
2033        // Check what contexts are included in each message
2034        let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
2035            (
2036                thread.message(message1_id).unwrap().clone(),
2037                thread.message(message2_id).unwrap().clone(),
2038                thread.message(message3_id).unwrap().clone(),
2039            )
2040        });
2041
2042        // First message should include context 1
2043        assert!(message1.context.contains("file1.rs"));
2044
2045        // Second message should include only context 2 (not 1)
2046        assert!(!message2.context.contains("file1.rs"));
2047        assert!(message2.context.contains("file2.rs"));
2048
2049        // Third message should include only context 3 (not 1 or 2)
2050        assert!(!message3.context.contains("file1.rs"));
2051        assert!(!message3.context.contains("file2.rs"));
2052        assert!(message3.context.contains("file3.rs"));
2053
2054        // Check entire request to make sure all contexts are properly included
2055        let request = thread.read_with(cx, |thread, cx| {
2056            thread.to_completion_request(RequestKind::Chat, cx)
2057        });
2058
2059        // The request should contain all 3 messages
2060        assert_eq!(request.messages.len(), 3);
2061
2062        // Check that the contexts are properly formatted in each message
2063        assert!(request.messages[0].string_contents().contains("file1.rs"));
2064        assert!(!request.messages[0].string_contents().contains("file2.rs"));
2065        assert!(!request.messages[0].string_contents().contains("file3.rs"));
2066
2067        assert!(!request.messages[1].string_contents().contains("file1.rs"));
2068        assert!(request.messages[1].string_contents().contains("file2.rs"));
2069        assert!(!request.messages[1].string_contents().contains("file3.rs"));
2070
2071        assert!(!request.messages[2].string_contents().contains("file1.rs"));
2072        assert!(!request.messages[2].string_contents().contains("file2.rs"));
2073        assert!(request.messages[2].string_contents().contains("file3.rs"));
2074    }
2075
2076    #[gpui::test]
2077    async fn test_message_without_files(cx: &mut TestAppContext) {
2078        init_test_settings(cx);
2079
2080        let project = create_test_project(
2081            cx,
2082            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
2083        )
2084        .await;
2085
2086        let (_, _thread_store, thread, _context_store) =
2087            setup_test_environment(cx, project.clone()).await;
2088
2089        // Insert user message without any context (empty context vector)
2090        let message_id = thread.update(cx, |thread, cx| {
2091            thread.insert_user_message("What is the best way to learn Rust?", vec![], None, cx)
2092        });
2093
2094        // Check content and context in message object
2095        let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
2096
2097        // Context should be empty when no files are included
2098        assert_eq!(message.role, Role::User);
2099        assert_eq!(message.segments.len(), 1);
2100        assert_eq!(
2101            message.segments[0],
2102            MessageSegment::Text("What is the best way to learn Rust?".to_string())
2103        );
2104        assert_eq!(message.context, "");
2105
2106        // Check message in request
2107        let request = thread.read_with(cx, |thread, cx| {
2108            thread.to_completion_request(RequestKind::Chat, cx)
2109        });
2110
2111        assert_eq!(request.messages.len(), 1);
2112        assert_eq!(
2113            request.messages[0].string_contents(),
2114            "What is the best way to learn Rust?"
2115        );
2116
2117        // Add second message, also without context
2118        let message2_id = thread.update(cx, |thread, cx| {
2119            thread.insert_user_message("Are there any good books?", vec![], None, cx)
2120        });
2121
2122        let message2 =
2123            thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
2124        assert_eq!(message2.context, "");
2125
2126        // Check that both messages appear in the request
2127        let request = thread.read_with(cx, |thread, cx| {
2128            thread.to_completion_request(RequestKind::Chat, cx)
2129        });
2130
2131        assert_eq!(request.messages.len(), 2);
2132        assert_eq!(
2133            request.messages[0].string_contents(),
2134            "What is the best way to learn Rust?"
2135        );
2136        assert_eq!(
2137            request.messages[1].string_contents(),
2138            "Are there any good books?"
2139        );
2140    }
2141
2142    #[gpui::test]
2143    async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
2144        init_test_settings(cx);
2145
2146        let project = create_test_project(
2147            cx,
2148            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
2149        )
2150        .await;
2151
2152        let (_workspace, _thread_store, thread, context_store) =
2153            setup_test_environment(cx, project.clone()).await;
2154
2155        // Open buffer and add it to context
2156        let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
2157            .await
2158            .unwrap();
2159
2160        let context =
2161            context_store.update(cx, |store, _| store.context().first().cloned().unwrap());
2162
2163        // Insert user message with the buffer as context
2164        thread.update(cx, |thread, cx| {
2165            thread.insert_user_message("Explain this code", vec![context], None, cx)
2166        });
2167
2168        // Create a request and check that it doesn't have a stale buffer warning yet
2169        let initial_request = thread.read_with(cx, |thread, cx| {
2170            thread.to_completion_request(RequestKind::Chat, cx)
2171        });
2172
2173        // Make sure we don't have a stale file warning yet
2174        let has_stale_warning = initial_request.messages.iter().any(|msg| {
2175            msg.string_contents()
2176                .contains("These files changed since last read:")
2177        });
2178        assert!(
2179            !has_stale_warning,
2180            "Should not have stale buffer warning before buffer is modified"
2181        );
2182
2183        // Modify the buffer
2184        buffer.update(cx, |buffer, cx| {
2185            // Find a position at the end of line 1
2186            buffer.edit(
2187                [(1..1, "\n    println!(\"Added a new line\");\n")],
2188                None,
2189                cx,
2190            );
2191        });
2192
2193        // Insert another user message without context
2194        thread.update(cx, |thread, cx| {
2195            thread.insert_user_message("What does the code do now?", vec![], None, cx)
2196        });
2197
2198        // Create a new request and check for the stale buffer warning
2199        let new_request = thread.read_with(cx, |thread, cx| {
2200            thread.to_completion_request(RequestKind::Chat, cx)
2201        });
2202
2203        // We should have a stale file warning as the last message
2204        let last_message = new_request
2205            .messages
2206            .last()
2207            .expect("Request should have messages");
2208
2209        // The last message should be the stale buffer notification
2210        assert_eq!(last_message.role, Role::User);
2211
2212        // Check the exact content of the message
2213        let expected_content = "These files changed since last read:\n- code.rs\n";
2214        assert_eq!(
2215            last_message.string_contents(),
2216            expected_content,
2217            "Last message should be exactly the stale buffer notification"
2218        );
2219    }
2220
2221    fn init_test_settings(cx: &mut TestAppContext) {
2222        cx.update(|cx| {
2223            let settings_store = SettingsStore::test(cx);
2224            cx.set_global(settings_store);
2225            language::init(cx);
2226            Project::init_settings(cx);
2227            AssistantSettings::register(cx);
2228            thread_store::init(cx);
2229            workspace::init_settings(cx);
2230            ThemeSettings::register(cx);
2231            ContextServerSettings::register(cx);
2232            EditorSettings::register(cx);
2233        });
2234    }
2235
2236    // Helper to create a test project with test files
2237    async fn create_test_project(
2238        cx: &mut TestAppContext,
2239        files: serde_json::Value,
2240    ) -> Entity<Project> {
2241        let fs = FakeFs::new(cx.executor());
2242        fs.insert_tree(path!("/test"), files).await;
2243        Project::test(fs, [path!("/test").as_ref()], cx).await
2244    }
2245
2246    async fn setup_test_environment(
2247        cx: &mut TestAppContext,
2248        project: Entity<Project>,
2249    ) -> (
2250        Entity<Workspace>,
2251        Entity<ThreadStore>,
2252        Entity<Thread>,
2253        Entity<ContextStore>,
2254    ) {
2255        let (workspace, cx) =
2256            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
2257
2258        let thread_store = cx.update(|_, cx| {
2259            ThreadStore::new(
2260                project.clone(),
2261                Arc::default(),
2262                Arc::new(PromptBuilder::new(None).unwrap()),
2263                cx,
2264            )
2265            .unwrap()
2266        });
2267
2268        let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
2269        let context_store = cx.new(|_cx| ContextStore::new(workspace.downgrade(), None));
2270
2271        (workspace, thread_store, thread, context_store)
2272    }
2273
2274    async fn add_file_to_context(
2275        project: &Entity<Project>,
2276        context_store: &Entity<ContextStore>,
2277        path: &str,
2278        cx: &mut TestAppContext,
2279    ) -> Result<Entity<language::Buffer>> {
2280        let buffer_path = project
2281            .read_with(cx, |project, cx| project.find_project_path(path, cx))
2282            .unwrap();
2283
2284        let buffer = project
2285            .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
2286            .await
2287            .unwrap();
2288
2289        context_store
2290            .update(cx, |store, cx| {
2291                store.add_file_from_buffer(buffer.clone(), cx)
2292            })
2293            .await?;
2294
2295        Ok(buffer)
2296    }
2297}