1use crate::{
   2    agent_profile::AgentProfile,
   3    context::{AgentContext, AgentContextHandle, ContextLoadResult, LoadedContext},
   4    thread_store::{
   5        SerializedCrease, SerializedLanguageModel, SerializedMessage, SerializedMessageSegment,
   6        SerializedThread, SerializedToolResult, SerializedToolUse, SharedProjectContext,
   7        ThreadStore,
   8    },
   9    tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState},
  10};
  11use agent_settings::{AgentProfileId, AgentSettings, CompletionMode};
  12use anyhow::{Result, anyhow};
  13use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
  14use chrono::{DateTime, Utc};
  15use client::{ModelRequestUsage, RequestUsage};
  16use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, Plan, UsageLimit};
  17use collections::HashMap;
  18use feature_flags::{self, FeatureFlagAppExt};
  19use futures::{FutureExt, StreamExt as _, future::Shared};
  20use git::repository::DiffType;
  21use gpui::{
  22    AnyWindowHandle, App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task,
  23    WeakEntity, Window,
  24};
  25use http_client::StatusCode;
  26use language_model::{
  27    ConfiguredModel, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
  28    LanguageModelExt as _, LanguageModelId, LanguageModelRegistry, LanguageModelRequest,
  29    LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
  30    LanguageModelToolResultContent, LanguageModelToolUse, LanguageModelToolUseId, MessageContent,
  31    ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, StopReason,
  32    TokenUsage,
  33};
  34use postage::stream::Stream as _;
  35use project::{
  36    Project,
  37    git_store::{GitStore, GitStoreCheckpoint, RepositoryState},
  38};
  39use prompt_store::{ModelContext, PromptBuilder};
  40use schemars::JsonSchema;
  41use serde::{Deserialize, Serialize};
  42use settings::Settings;
  43use std::{
  44    io::Write,
  45    ops::Range,
  46    sync::Arc,
  47    time::{Duration, Instant},
  48};
  49use thiserror::Error;
  50use util::{ResultExt as _, post_inc};
  51use uuid::Uuid;
  52
  53const MAX_RETRY_ATTEMPTS: u8 = 4;
  54const BASE_RETRY_DELAY: Duration = Duration::from_secs(5);
  55
  56#[derive(Debug, Clone)]
  57enum RetryStrategy {
  58    ExponentialBackoff {
  59        initial_delay: Duration,
  60        max_attempts: u8,
  61    },
  62    Fixed {
  63        delay: Duration,
  64        max_attempts: u8,
  65    },
  66}
  67
  68#[derive(
  69    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
  70)]
  71pub struct ThreadId(Arc<str>);
  72
  73impl ThreadId {
  74    pub fn new() -> Self {
  75        Self(Uuid::new_v4().to_string().into())
  76    }
  77}
  78
  79impl std::fmt::Display for ThreadId {
  80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  81        write!(f, "{}", self.0)
  82    }
  83}
  84
  85impl From<&str> for ThreadId {
  86    fn from(value: &str) -> Self {
  87        Self(value.into())
  88    }
  89}
  90
  91/// The ID of the user prompt that initiated a request.
  92///
  93/// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).
  94#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
  95pub struct PromptId(Arc<str>);
  96
  97impl PromptId {
  98    pub fn new() -> Self {
  99        Self(Uuid::new_v4().to_string().into())
 100    }
 101}
 102
 103impl std::fmt::Display for PromptId {
 104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 105        write!(f, "{}", self.0)
 106    }
 107}
 108
 109#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
 110pub struct MessageId(pub(crate) usize);
 111
 112impl MessageId {
 113    fn post_inc(&mut self) -> Self {
 114        Self(post_inc(&mut self.0))
 115    }
 116
 117    pub fn as_usize(&self) -> usize {
 118        self.0
 119    }
 120}
 121
 122/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
 123#[derive(Clone, Debug)]
 124pub struct MessageCrease {
 125    pub range: Range<usize>,
 126    pub icon_path: SharedString,
 127    pub label: SharedString,
 128    /// None for a deserialized message, Some otherwise.
 129    pub context: Option<AgentContextHandle>,
 130}
 131
 132/// A message in a [`Thread`].
 133#[derive(Debug, Clone)]
 134pub struct Message {
 135    pub id: MessageId,
 136    pub role: Role,
 137    pub segments: Vec<MessageSegment>,
 138    pub loaded_context: LoadedContext,
 139    pub creases: Vec<MessageCrease>,
 140    pub is_hidden: bool,
 141    pub ui_only: bool,
 142}
 143
 144impl Message {
 145    /// Returns whether the message contains any meaningful text that should be displayed
 146    /// The model sometimes runs tool without producing any text or just a marker ([`USING_TOOL_MARKER`])
 147    pub fn should_display_content(&self) -> bool {
 148        self.segments.iter().all(|segment| segment.should_display())
 149    }
 150
 151    pub fn push_thinking(&mut self, text: &str, signature: Option<String>) {
 152        if let Some(MessageSegment::Thinking {
 153            text: segment,
 154            signature: current_signature,
 155        }) = self.segments.last_mut()
 156        {
 157            if let Some(signature) = signature {
 158                *current_signature = Some(signature);
 159            }
 160            segment.push_str(text);
 161        } else {
 162            self.segments.push(MessageSegment::Thinking {
 163                text: text.to_string(),
 164                signature,
 165            });
 166        }
 167    }
 168
 169    pub fn push_redacted_thinking(&mut self, data: String) {
 170        self.segments.push(MessageSegment::RedactedThinking(data));
 171    }
 172
 173    pub fn push_text(&mut self, text: &str) {
 174        if let Some(MessageSegment::Text(segment)) = self.segments.last_mut() {
 175            segment.push_str(text);
 176        } else {
 177            self.segments.push(MessageSegment::Text(text.to_string()));
 178        }
 179    }
 180
 181    pub fn to_string(&self) -> String {
 182        let mut result = String::new();
 183
 184        if !self.loaded_context.text.is_empty() {
 185            result.push_str(&self.loaded_context.text);
 186        }
 187
 188        for segment in &self.segments {
 189            match segment {
 190                MessageSegment::Text(text) => result.push_str(text),
 191                MessageSegment::Thinking { text, .. } => {
 192                    result.push_str("<think>\n");
 193                    result.push_str(text);
 194                    result.push_str("\n</think>");
 195                }
 196                MessageSegment::RedactedThinking(_) => {}
 197            }
 198        }
 199
 200        result
 201    }
 202}
 203
 204#[derive(Debug, Clone, PartialEq, Eq)]
 205pub enum MessageSegment {
 206    Text(String),
 207    Thinking {
 208        text: String,
 209        signature: Option<String>,
 210    },
 211    RedactedThinking(String),
 212}
 213
 214impl MessageSegment {
 215    pub fn should_display(&self) -> bool {
 216        match self {
 217            Self::Text(text) => text.is_empty(),
 218            Self::Thinking { text, .. } => text.is_empty(),
 219            Self::RedactedThinking(_) => false,
 220        }
 221    }
 222
 223    pub fn text(&self) -> Option<&str> {
 224        match self {
 225            MessageSegment::Text(text) => Some(text),
 226            _ => None,
 227        }
 228    }
 229}
 230
 231#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 232pub struct ProjectSnapshot {
 233    pub worktree_snapshots: Vec<WorktreeSnapshot>,
 234    pub unsaved_buffer_paths: Vec<String>,
 235    pub timestamp: DateTime<Utc>,
 236}
 237
 238#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 239pub struct WorktreeSnapshot {
 240    pub worktree_path: String,
 241    pub git_state: Option<GitState>,
 242}
 243
 244#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 245pub struct GitState {
 246    pub remote_url: Option<String>,
 247    pub head_sha: Option<String>,
 248    pub current_branch: Option<String>,
 249    pub diff: Option<String>,
 250}
 251
 252#[derive(Clone, Debug)]
 253pub struct ThreadCheckpoint {
 254    message_id: MessageId,
 255    git_checkpoint: GitStoreCheckpoint,
 256}
 257
 258#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 259pub enum ThreadFeedback {
 260    Positive,
 261    Negative,
 262}
 263
 264pub enum LastRestoreCheckpoint {
 265    Pending {
 266        message_id: MessageId,
 267    },
 268    Error {
 269        message_id: MessageId,
 270        error: String,
 271    },
 272}
 273
 274impl LastRestoreCheckpoint {
 275    pub fn message_id(&self) -> MessageId {
 276        match self {
 277            LastRestoreCheckpoint::Pending { message_id } => *message_id,
 278            LastRestoreCheckpoint::Error { message_id, .. } => *message_id,
 279        }
 280    }
 281}
 282
 283#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
 284pub enum DetailedSummaryState {
 285    #[default]
 286    NotGenerated,
 287    Generating {
 288        message_id: MessageId,
 289    },
 290    Generated {
 291        text: SharedString,
 292        message_id: MessageId,
 293    },
 294}
 295
 296impl DetailedSummaryState {
 297    fn text(&self) -> Option<SharedString> {
 298        if let Self::Generated { text, .. } = self {
 299            Some(text.clone())
 300        } else {
 301            None
 302        }
 303    }
 304}
 305
 306#[derive(Default, Debug)]
 307pub struct TotalTokenUsage {
 308    pub total: u64,
 309    pub max: u64,
 310}
 311
 312impl TotalTokenUsage {
 313    pub fn ratio(&self) -> TokenUsageRatio {
 314        #[cfg(debug_assertions)]
 315        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
 316            .unwrap_or("0.8".to_string())
 317            .parse()
 318            .unwrap();
 319        #[cfg(not(debug_assertions))]
 320        let warning_threshold: f32 = 0.8;
 321
 322        // When the maximum is unknown because there is no selected model,
 323        // avoid showing the token limit warning.
 324        if self.max == 0 {
 325            TokenUsageRatio::Normal
 326        } else if self.total >= self.max {
 327            TokenUsageRatio::Exceeded
 328        } else if self.total as f32 / self.max as f32 >= warning_threshold {
 329            TokenUsageRatio::Warning
 330        } else {
 331            TokenUsageRatio::Normal
 332        }
 333    }
 334
 335    pub fn add(&self, tokens: u64) -> TotalTokenUsage {
 336        TotalTokenUsage {
 337            total: self.total + tokens,
 338            max: self.max,
 339        }
 340    }
 341}
 342
 343#[derive(Debug, Default, PartialEq, Eq)]
 344pub enum TokenUsageRatio {
 345    #[default]
 346    Normal,
 347    Warning,
 348    Exceeded,
 349}
 350
 351#[derive(Debug, Clone, Copy)]
 352pub enum QueueState {
 353    Sending,
 354    Queued { position: usize },
 355    Started,
 356}
 357
 358/// A thread of conversation with the LLM.
 359pub struct Thread {
 360    id: ThreadId,
 361    updated_at: DateTime<Utc>,
 362    summary: ThreadSummary,
 363    pending_summary: Task<Option<()>>,
 364    detailed_summary_task: Task<Option<()>>,
 365    detailed_summary_tx: postage::watch::Sender<DetailedSummaryState>,
 366    detailed_summary_rx: postage::watch::Receiver<DetailedSummaryState>,
 367    completion_mode: agent_settings::CompletionMode,
 368    messages: Vec<Message>,
 369    next_message_id: MessageId,
 370    last_prompt_id: PromptId,
 371    project_context: SharedProjectContext,
 372    checkpoints_by_message: HashMap<MessageId, ThreadCheckpoint>,
 373    completion_count: usize,
 374    pending_completions: Vec<PendingCompletion>,
 375    project: Entity<Project>,
 376    prompt_builder: Arc<PromptBuilder>,
 377    tools: Entity<ToolWorkingSet>,
 378    tool_use: ToolUseState,
 379    action_log: Entity<ActionLog>,
 380    last_restore_checkpoint: Option<LastRestoreCheckpoint>,
 381    pending_checkpoint: Option<ThreadCheckpoint>,
 382    initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
 383    request_token_usage: Vec<TokenUsage>,
 384    cumulative_token_usage: TokenUsage,
 385    exceeded_window_error: Option<ExceededWindowError>,
 386    tool_use_limit_reached: bool,
 387    feedback: Option<ThreadFeedback>,
 388    retry_state: Option<RetryState>,
 389    message_feedback: HashMap<MessageId, ThreadFeedback>,
 390    last_auto_capture_at: Option<Instant>,
 391    last_received_chunk_at: Option<Instant>,
 392    request_callback: Option<
 393        Box<dyn FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>])>,
 394    >,
 395    remaining_turns: u32,
 396    configured_model: Option<ConfiguredModel>,
 397    profile: AgentProfile,
 398    last_error_context: Option<(Arc<dyn LanguageModel>, CompletionIntent)>,
 399}
 400
 401#[derive(Clone, Debug)]
 402struct RetryState {
 403    attempt: u8,
 404    max_attempts: u8,
 405    intent: CompletionIntent,
 406}
 407
 408#[derive(Clone, Debug, PartialEq, Eq)]
 409pub enum ThreadSummary {
 410    Pending,
 411    Generating,
 412    Ready(SharedString),
 413    Error,
 414}
 415
 416impl ThreadSummary {
 417    pub const DEFAULT: SharedString = SharedString::new_static("New Thread");
 418
 419    pub fn or_default(&self) -> SharedString {
 420        self.unwrap_or(Self::DEFAULT)
 421    }
 422
 423    pub fn unwrap_or(&self, message: impl Into<SharedString>) -> SharedString {
 424        self.ready().unwrap_or_else(|| message.into())
 425    }
 426
 427    pub fn ready(&self) -> Option<SharedString> {
 428        match self {
 429            ThreadSummary::Ready(summary) => Some(summary.clone()),
 430            ThreadSummary::Pending | ThreadSummary::Generating | ThreadSummary::Error => None,
 431        }
 432    }
 433}
 434
 435#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 436pub struct ExceededWindowError {
 437    /// Model used when last message exceeded context window
 438    model_id: LanguageModelId,
 439    /// Token count including last message
 440    token_count: u64,
 441}
 442
 443impl Thread {
 444    pub fn new(
 445        project: Entity<Project>,
 446        tools: Entity<ToolWorkingSet>,
 447        prompt_builder: Arc<PromptBuilder>,
 448        system_prompt: SharedProjectContext,
 449        cx: &mut Context<Self>,
 450    ) -> Self {
 451        let (detailed_summary_tx, detailed_summary_rx) = postage::watch::channel();
 452        let configured_model = LanguageModelRegistry::read_global(cx).default_model();
 453        let profile_id = AgentSettings::get_global(cx).default_profile.clone();
 454
 455        Self {
 456            id: ThreadId::new(),
 457            updated_at: Utc::now(),
 458            summary: ThreadSummary::Pending,
 459            pending_summary: Task::ready(None),
 460            detailed_summary_task: Task::ready(None),
 461            detailed_summary_tx,
 462            detailed_summary_rx,
 463            completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
 464            messages: Vec::new(),
 465            next_message_id: MessageId(0),
 466            last_prompt_id: PromptId::new(),
 467            project_context: system_prompt,
 468            checkpoints_by_message: HashMap::default(),
 469            completion_count: 0,
 470            pending_completions: Vec::new(),
 471            project: project.clone(),
 472            prompt_builder,
 473            tools: tools.clone(),
 474            last_restore_checkpoint: None,
 475            pending_checkpoint: None,
 476            tool_use: ToolUseState::new(tools.clone()),
 477            action_log: cx.new(|_| ActionLog::new(project.clone())),
 478            initial_project_snapshot: {
 479                let project_snapshot = Self::project_snapshot(project, cx);
 480                cx.foreground_executor()
 481                    .spawn(async move { Some(project_snapshot.await) })
 482                    .shared()
 483            },
 484            request_token_usage: Vec::new(),
 485            cumulative_token_usage: TokenUsage::default(),
 486            exceeded_window_error: None,
 487            tool_use_limit_reached: false,
 488            feedback: None,
 489            retry_state: None,
 490            message_feedback: HashMap::default(),
 491            last_auto_capture_at: None,
 492            last_error_context: None,
 493            last_received_chunk_at: None,
 494            request_callback: None,
 495            remaining_turns: u32::MAX,
 496            configured_model: configured_model.clone(),
 497            profile: AgentProfile::new(profile_id, tools),
 498        }
 499    }
 500
 501    pub fn deserialize(
 502        id: ThreadId,
 503        serialized: SerializedThread,
 504        project: Entity<Project>,
 505        tools: Entity<ToolWorkingSet>,
 506        prompt_builder: Arc<PromptBuilder>,
 507        project_context: SharedProjectContext,
 508        window: Option<&mut Window>, // None in headless mode
 509        cx: &mut Context<Self>,
 510    ) -> Self {
 511        let next_message_id = MessageId(
 512            serialized
 513                .messages
 514                .last()
 515                .map(|message| message.id.0 + 1)
 516                .unwrap_or(0),
 517        );
 518        let tool_use = ToolUseState::from_serialized_messages(
 519            tools.clone(),
 520            &serialized.messages,
 521            project.clone(),
 522            window,
 523            cx,
 524        );
 525        let (detailed_summary_tx, detailed_summary_rx) =
 526            postage::watch::channel_with(serialized.detailed_summary_state);
 527
 528        let configured_model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
 529            serialized
 530                .model
 531                .and_then(|model| {
 532                    let model = SelectedModel {
 533                        provider: model.provider.clone().into(),
 534                        model: model.model.clone().into(),
 535                    };
 536                    registry.select_model(&model, cx)
 537                })
 538                .or_else(|| registry.default_model())
 539        });
 540
 541        let completion_mode = serialized
 542            .completion_mode
 543            .unwrap_or_else(|| AgentSettings::get_global(cx).preferred_completion_mode);
 544        let profile_id = serialized
 545            .profile
 546            .unwrap_or_else(|| AgentSettings::get_global(cx).default_profile.clone());
 547
 548        Self {
 549            id,
 550            updated_at: serialized.updated_at,
 551            summary: ThreadSummary::Ready(serialized.summary),
 552            pending_summary: Task::ready(None),
 553            detailed_summary_task: Task::ready(None),
 554            detailed_summary_tx,
 555            detailed_summary_rx,
 556            completion_mode,
 557            retry_state: None,
 558            messages: serialized
 559                .messages
 560                .into_iter()
 561                .map(|message| Message {
 562                    id: message.id,
 563                    role: message.role,
 564                    segments: message
 565                        .segments
 566                        .into_iter()
 567                        .map(|segment| match segment {
 568                            SerializedMessageSegment::Text { text } => MessageSegment::Text(text),
 569                            SerializedMessageSegment::Thinking { text, signature } => {
 570                                MessageSegment::Thinking { text, signature }
 571                            }
 572                            SerializedMessageSegment::RedactedThinking { data } => {
 573                                MessageSegment::RedactedThinking(data)
 574                            }
 575                        })
 576                        .collect(),
 577                    loaded_context: LoadedContext {
 578                        contexts: Vec::new(),
 579                        text: message.context,
 580                        images: Vec::new(),
 581                    },
 582                    creases: message
 583                        .creases
 584                        .into_iter()
 585                        .map(|crease| MessageCrease {
 586                            range: crease.start..crease.end,
 587                            icon_path: crease.icon_path,
 588                            label: crease.label,
 589                            context: None,
 590                        })
 591                        .collect(),
 592                    is_hidden: message.is_hidden,
 593                    ui_only: false, // UI-only messages are not persisted
 594                })
 595                .collect(),
 596            next_message_id,
 597            last_prompt_id: PromptId::new(),
 598            project_context,
 599            checkpoints_by_message: HashMap::default(),
 600            completion_count: 0,
 601            pending_completions: Vec::new(),
 602            last_restore_checkpoint: None,
 603            pending_checkpoint: None,
 604            project: project.clone(),
 605            prompt_builder,
 606            tools: tools.clone(),
 607            tool_use,
 608            action_log: cx.new(|_| ActionLog::new(project)),
 609            initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
 610            request_token_usage: serialized.request_token_usage,
 611            cumulative_token_usage: serialized.cumulative_token_usage,
 612            exceeded_window_error: None,
 613            tool_use_limit_reached: serialized.tool_use_limit_reached,
 614            feedback: None,
 615            message_feedback: HashMap::default(),
 616            last_auto_capture_at: None,
 617            last_error_context: None,
 618            last_received_chunk_at: None,
 619            request_callback: None,
 620            remaining_turns: u32::MAX,
 621            configured_model,
 622            profile: AgentProfile::new(profile_id, tools),
 623        }
 624    }
 625
 626    pub fn set_request_callback(
 627        &mut self,
 628        callback: impl 'static
 629        + FnMut(&LanguageModelRequest, &[Result<LanguageModelCompletionEvent, String>]),
 630    ) {
 631        self.request_callback = Some(Box::new(callback));
 632    }
 633
 634    pub fn id(&self) -> &ThreadId {
 635        &self.id
 636    }
 637
 638    pub fn profile(&self) -> &AgentProfile {
 639        &self.profile
 640    }
 641
 642    pub fn set_profile(&mut self, id: AgentProfileId, cx: &mut Context<Self>) {
 643        if &id != self.profile.id() {
 644            self.profile = AgentProfile::new(id, self.tools.clone());
 645            cx.emit(ThreadEvent::ProfileChanged);
 646        }
 647    }
 648
 649    pub fn is_empty(&self) -> bool {
 650        self.messages.is_empty()
 651    }
 652
 653    pub fn updated_at(&self) -> DateTime<Utc> {
 654        self.updated_at
 655    }
 656
 657    pub fn touch_updated_at(&mut self) {
 658        self.updated_at = Utc::now();
 659    }
 660
 661    pub fn advance_prompt_id(&mut self) {
 662        self.last_prompt_id = PromptId::new();
 663    }
 664
 665    pub fn project_context(&self) -> SharedProjectContext {
 666        self.project_context.clone()
 667    }
 668
 669    pub fn get_or_init_configured_model(&mut self, cx: &App) -> Option<ConfiguredModel> {
 670        if self.configured_model.is_none() {
 671            self.configured_model = LanguageModelRegistry::read_global(cx).default_model();
 672        }
 673        self.configured_model.clone()
 674    }
 675
 676    pub fn configured_model(&self) -> Option<ConfiguredModel> {
 677        self.configured_model.clone()
 678    }
 679
 680    pub fn set_configured_model(&mut self, model: Option<ConfiguredModel>, cx: &mut Context<Self>) {
 681        self.configured_model = model;
 682        cx.notify();
 683    }
 684
 685    pub fn summary(&self) -> &ThreadSummary {
 686        &self.summary
 687    }
 688
 689    pub fn set_summary(&mut self, new_summary: impl Into<SharedString>, cx: &mut Context<Self>) {
 690        let current_summary = match &self.summary {
 691            ThreadSummary::Pending | ThreadSummary::Generating => return,
 692            ThreadSummary::Ready(summary) => summary,
 693            ThreadSummary::Error => &ThreadSummary::DEFAULT,
 694        };
 695
 696        let mut new_summary = new_summary.into();
 697
 698        if new_summary.is_empty() {
 699            new_summary = ThreadSummary::DEFAULT;
 700        }
 701
 702        if current_summary != &new_summary {
 703            self.summary = ThreadSummary::Ready(new_summary);
 704            cx.emit(ThreadEvent::SummaryChanged);
 705        }
 706    }
 707
 708    pub fn completion_mode(&self) -> CompletionMode {
 709        self.completion_mode
 710    }
 711
 712    pub fn set_completion_mode(&mut self, mode: CompletionMode) {
 713        self.completion_mode = mode;
 714    }
 715
 716    pub fn message(&self, id: MessageId) -> Option<&Message> {
 717        let index = self
 718            .messages
 719            .binary_search_by(|message| message.id.cmp(&id))
 720            .ok()?;
 721
 722        self.messages.get(index)
 723    }
 724
 725    pub fn messages(&self) -> impl ExactSizeIterator<Item = &Message> {
 726        self.messages.iter()
 727    }
 728
 729    pub fn is_generating(&self) -> bool {
 730        !self.pending_completions.is_empty() || !self.all_tools_finished()
 731    }
 732
 733    /// Indicates whether streaming of language model events is stale.
 734    /// When `is_generating()` is false, this method returns `None`.
 735    pub fn is_generation_stale(&self) -> Option<bool> {
 736        const STALE_THRESHOLD: u128 = 250;
 737
 738        self.last_received_chunk_at
 739            .map(|instant| instant.elapsed().as_millis() > STALE_THRESHOLD)
 740    }
 741
 742    fn received_chunk(&mut self) {
 743        self.last_received_chunk_at = Some(Instant::now());
 744    }
 745
 746    pub fn queue_state(&self) -> Option<QueueState> {
 747        self.pending_completions
 748            .first()
 749            .map(|pending_completion| pending_completion.queue_state)
 750    }
 751
 752    pub fn tools(&self) -> &Entity<ToolWorkingSet> {
 753        &self.tools
 754    }
 755
 756    pub fn pending_tool(&self, id: &LanguageModelToolUseId) -> Option<&PendingToolUse> {
 757        self.tool_use
 758            .pending_tool_uses()
 759            .into_iter()
 760            .find(|tool_use| &tool_use.id == id)
 761    }
 762
 763    pub fn tools_needing_confirmation(&self) -> impl Iterator<Item = &PendingToolUse> {
 764        self.tool_use
 765            .pending_tool_uses()
 766            .into_iter()
 767            .filter(|tool_use| tool_use.status.needs_confirmation())
 768    }
 769
 770    pub fn has_pending_tool_uses(&self) -> bool {
 771        !self.tool_use.pending_tool_uses().is_empty()
 772    }
 773
 774    pub fn checkpoint_for_message(&self, id: MessageId) -> Option<ThreadCheckpoint> {
 775        self.checkpoints_by_message.get(&id).cloned()
 776    }
 777
 778    pub fn restore_checkpoint(
 779        &mut self,
 780        checkpoint: ThreadCheckpoint,
 781        cx: &mut Context<Self>,
 782    ) -> Task<Result<()>> {
 783        self.last_restore_checkpoint = Some(LastRestoreCheckpoint::Pending {
 784            message_id: checkpoint.message_id,
 785        });
 786        cx.emit(ThreadEvent::CheckpointChanged);
 787        cx.notify();
 788
 789        let git_store = self.project().read(cx).git_store().clone();
 790        let restore = git_store.update(cx, |git_store, cx| {
 791            git_store.restore_checkpoint(checkpoint.git_checkpoint.clone(), cx)
 792        });
 793
 794        cx.spawn(async move |this, cx| {
 795            let result = restore.await;
 796            this.update(cx, |this, cx| {
 797                if let Err(err) = result.as_ref() {
 798                    this.last_restore_checkpoint = Some(LastRestoreCheckpoint::Error {
 799                        message_id: checkpoint.message_id,
 800                        error: err.to_string(),
 801                    });
 802                } else {
 803                    this.truncate(checkpoint.message_id, cx);
 804                    this.last_restore_checkpoint = None;
 805                }
 806                this.pending_checkpoint = None;
 807                cx.emit(ThreadEvent::CheckpointChanged);
 808                cx.notify();
 809            })?;
 810            result
 811        })
 812    }
 813
 814    fn finalize_pending_checkpoint(&mut self, cx: &mut Context<Self>) {
 815        let pending_checkpoint = if self.is_generating() {
 816            return;
 817        } else if let Some(checkpoint) = self.pending_checkpoint.take() {
 818            checkpoint
 819        } else {
 820            return;
 821        };
 822
 823        self.finalize_checkpoint(pending_checkpoint, cx);
 824    }
 825
 826    fn finalize_checkpoint(
 827        &mut self,
 828        pending_checkpoint: ThreadCheckpoint,
 829        cx: &mut Context<Self>,
 830    ) {
 831        let git_store = self.project.read(cx).git_store().clone();
 832        let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
 833        cx.spawn(async move |this, cx| match final_checkpoint.await {
 834            Ok(final_checkpoint) => {
 835                let equal = git_store
 836                    .update(cx, |store, cx| {
 837                        store.compare_checkpoints(
 838                            pending_checkpoint.git_checkpoint.clone(),
 839                            final_checkpoint.clone(),
 840                            cx,
 841                        )
 842                    })?
 843                    .await
 844                    .unwrap_or(false);
 845
 846                if !equal {
 847                    this.update(cx, |this, cx| {
 848                        this.insert_checkpoint(pending_checkpoint, cx)
 849                    })?;
 850                }
 851
 852                Ok(())
 853            }
 854            Err(_) => this.update(cx, |this, cx| {
 855                this.insert_checkpoint(pending_checkpoint, cx)
 856            }),
 857        })
 858        .detach();
 859    }
 860
 861    fn insert_checkpoint(&mut self, checkpoint: ThreadCheckpoint, cx: &mut Context<Self>) {
 862        self.checkpoints_by_message
 863            .insert(checkpoint.message_id, checkpoint);
 864        cx.emit(ThreadEvent::CheckpointChanged);
 865        cx.notify();
 866    }
 867
 868    pub fn last_restore_checkpoint(&self) -> Option<&LastRestoreCheckpoint> {
 869        self.last_restore_checkpoint.as_ref()
 870    }
 871
 872    pub fn truncate(&mut self, message_id: MessageId, cx: &mut Context<Self>) {
 873        let Some(message_ix) = self
 874            .messages
 875            .iter()
 876            .rposition(|message| message.id == message_id)
 877        else {
 878            return;
 879        };
 880        for deleted_message in self.messages.drain(message_ix..) {
 881            self.checkpoints_by_message.remove(&deleted_message.id);
 882        }
 883        cx.notify();
 884    }
 885
 886    pub fn context_for_message(&self, id: MessageId) -> impl Iterator<Item = &AgentContext> {
 887        self.messages
 888            .iter()
 889            .find(|message| message.id == id)
 890            .into_iter()
 891            .flat_map(|message| message.loaded_context.contexts.iter())
 892    }
 893
 894    pub fn is_turn_end(&self, ix: usize) -> bool {
 895        if self.messages.is_empty() {
 896            return false;
 897        }
 898
 899        if !self.is_generating() && ix == self.messages.len() - 1 {
 900            return true;
 901        }
 902
 903        let Some(message) = self.messages.get(ix) else {
 904            return false;
 905        };
 906
 907        if message.role != Role::Assistant {
 908            return false;
 909        }
 910
 911        self.messages
 912            .get(ix + 1)
 913            .and_then(|message| {
 914                self.message(message.id)
 915                    .map(|next_message| next_message.role == Role::User && !next_message.is_hidden)
 916            })
 917            .unwrap_or(false)
 918    }
 919
 920    pub fn tool_use_limit_reached(&self) -> bool {
 921        self.tool_use_limit_reached
 922    }
 923
 924    /// Returns whether all of the tool uses have finished running.
 925    pub fn all_tools_finished(&self) -> bool {
 926        // If the only pending tool uses left are the ones with errors, then
 927        // that means that we've finished running all of the pending tools.
 928        self.tool_use
 929            .pending_tool_uses()
 930            .iter()
 931            .all(|pending_tool_use| pending_tool_use.status.is_error())
 932    }
 933
 934    /// Returns whether any pending tool uses may perform edits
 935    pub fn has_pending_edit_tool_uses(&self) -> bool {
 936        self.tool_use
 937            .pending_tool_uses()
 938            .iter()
 939            .filter(|pending_tool_use| !pending_tool_use.status.is_error())
 940            .any(|pending_tool_use| pending_tool_use.may_perform_edits)
 941    }
 942
 943    pub fn tool_uses_for_message(&self, id: MessageId, cx: &App) -> Vec<ToolUse> {
 944        self.tool_use.tool_uses_for_message(id, &self.project, cx)
 945    }
 946
 947    pub fn tool_results_for_message(
 948        &self,
 949        assistant_message_id: MessageId,
 950    ) -> Vec<&LanguageModelToolResult> {
 951        self.tool_use.tool_results_for_message(assistant_message_id)
 952    }
 953
 954    pub fn tool_result(&self, id: &LanguageModelToolUseId) -> Option<&LanguageModelToolResult> {
 955        self.tool_use.tool_result(id)
 956    }
 957
 958    pub fn output_for_tool(&self, id: &LanguageModelToolUseId) -> Option<&Arc<str>> {
 959        match &self.tool_use.tool_result(id)?.content {
 960            LanguageModelToolResultContent::Text(text) => Some(text),
 961            LanguageModelToolResultContent::Image(_) => {
 962                // TODO: We should display image
 963                None
 964            }
 965        }
 966    }
 967
 968    pub fn card_for_tool(&self, id: &LanguageModelToolUseId) -> Option<AnyToolCard> {
 969        self.tool_use.tool_result_card(id).cloned()
 970    }
 971
 972    /// Return tools that are both enabled and supported by the model
 973    pub fn available_tools(
 974        &self,
 975        cx: &App,
 976        model: Arc<dyn LanguageModel>,
 977    ) -> Vec<LanguageModelRequestTool> {
 978        if model.supports_tools() {
 979            self.profile
 980                .enabled_tools(cx)
 981                .into_iter()
 982                .filter_map(|(name, tool)| {
 983                    // Skip tools that cannot be supported
 984                    let input_schema = tool.input_schema(model.tool_input_format()).ok()?;
 985                    Some(LanguageModelRequestTool {
 986                        name: name.into(),
 987                        description: tool.description(),
 988                        input_schema,
 989                    })
 990                })
 991                .collect()
 992        } else {
 993            Vec::default()
 994        }
 995    }
 996
 997    pub fn insert_user_message(
 998        &mut self,
 999        text: impl Into<String>,
1000        loaded_context: ContextLoadResult,
1001        git_checkpoint: Option<GitStoreCheckpoint>,
1002        creases: Vec<MessageCrease>,
1003        cx: &mut Context<Self>,
1004    ) -> MessageId {
1005        if !loaded_context.referenced_buffers.is_empty() {
1006            self.action_log.update(cx, |log, cx| {
1007                for buffer in loaded_context.referenced_buffers {
1008                    log.buffer_read(buffer, cx);
1009                }
1010            });
1011        }
1012
1013        let message_id = self.insert_message(
1014            Role::User,
1015            vec![MessageSegment::Text(text.into())],
1016            loaded_context.loaded_context,
1017            creases,
1018            false,
1019            cx,
1020        );
1021
1022        if let Some(git_checkpoint) = git_checkpoint {
1023            self.pending_checkpoint = Some(ThreadCheckpoint {
1024                message_id,
1025                git_checkpoint,
1026            });
1027        }
1028
1029        self.auto_capture_telemetry(cx);
1030
1031        message_id
1032    }
1033
1034    pub fn insert_invisible_continue_message(&mut self, cx: &mut Context<Self>) -> MessageId {
1035        let id = self.insert_message(
1036            Role::User,
1037            vec![MessageSegment::Text("Continue where you left off".into())],
1038            LoadedContext::default(),
1039            vec![],
1040            true,
1041            cx,
1042        );
1043        self.pending_checkpoint = None;
1044
1045        id
1046    }
1047
1048    pub fn insert_assistant_message(
1049        &mut self,
1050        segments: Vec<MessageSegment>,
1051        cx: &mut Context<Self>,
1052    ) -> MessageId {
1053        self.insert_message(
1054            Role::Assistant,
1055            segments,
1056            LoadedContext::default(),
1057            Vec::new(),
1058            false,
1059            cx,
1060        )
1061    }
1062
1063    pub fn insert_message(
1064        &mut self,
1065        role: Role,
1066        segments: Vec<MessageSegment>,
1067        loaded_context: LoadedContext,
1068        creases: Vec<MessageCrease>,
1069        is_hidden: bool,
1070        cx: &mut Context<Self>,
1071    ) -> MessageId {
1072        let id = self.next_message_id.post_inc();
1073        self.messages.push(Message {
1074            id,
1075            role,
1076            segments,
1077            loaded_context,
1078            creases,
1079            is_hidden,
1080            ui_only: false,
1081        });
1082        self.touch_updated_at();
1083        cx.emit(ThreadEvent::MessageAdded(id));
1084        id
1085    }
1086
1087    pub fn edit_message(
1088        &mut self,
1089        id: MessageId,
1090        new_role: Role,
1091        new_segments: Vec<MessageSegment>,
1092        creases: Vec<MessageCrease>,
1093        loaded_context: Option<LoadedContext>,
1094        checkpoint: Option<GitStoreCheckpoint>,
1095        cx: &mut Context<Self>,
1096    ) -> bool {
1097        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
1098            return false;
1099        };
1100        message.role = new_role;
1101        message.segments = new_segments;
1102        message.creases = creases;
1103        if let Some(context) = loaded_context {
1104            message.loaded_context = context;
1105        }
1106        if let Some(git_checkpoint) = checkpoint {
1107            self.checkpoints_by_message.insert(
1108                id,
1109                ThreadCheckpoint {
1110                    message_id: id,
1111                    git_checkpoint,
1112                },
1113            );
1114        }
1115        self.touch_updated_at();
1116        cx.emit(ThreadEvent::MessageEdited(id));
1117        true
1118    }
1119
1120    pub fn delete_message(&mut self, id: MessageId, cx: &mut Context<Self>) -> bool {
1121        let Some(index) = self.messages.iter().position(|message| message.id == id) else {
1122            return false;
1123        };
1124        self.messages.remove(index);
1125        self.touch_updated_at();
1126        cx.emit(ThreadEvent::MessageDeleted(id));
1127        true
1128    }
1129
1130    /// Returns the representation of this [`Thread`] in a textual form.
1131    ///
1132    /// This is the representation we use when attaching a thread as context to another thread.
1133    pub fn text(&self) -> String {
1134        let mut text = String::new();
1135
1136        for message in &self.messages {
1137            text.push_str(match message.role {
1138                language_model::Role::User => "User:",
1139                language_model::Role::Assistant => "Agent:",
1140                language_model::Role::System => "System:",
1141            });
1142            text.push('\n');
1143
1144            for segment in &message.segments {
1145                match segment {
1146                    MessageSegment::Text(content) => text.push_str(content),
1147                    MessageSegment::Thinking { text: content, .. } => {
1148                        text.push_str(&format!("<think>{}</think>", content))
1149                    }
1150                    MessageSegment::RedactedThinking(_) => {}
1151                }
1152            }
1153            text.push('\n');
1154        }
1155
1156        text
1157    }
1158
1159    /// Serializes this thread into a format for storage or telemetry.
1160    pub fn serialize(&self, cx: &mut Context<Self>) -> Task<Result<SerializedThread>> {
1161        let initial_project_snapshot = self.initial_project_snapshot.clone();
1162        cx.spawn(async move |this, cx| {
1163            let initial_project_snapshot = initial_project_snapshot.await;
1164            this.read_with(cx, |this, cx| SerializedThread {
1165                version: SerializedThread::VERSION.to_string(),
1166                summary: this.summary().or_default(),
1167                updated_at: this.updated_at(),
1168                messages: this
1169                    .messages()
1170                    .filter(|message| !message.ui_only)
1171                    .map(|message| SerializedMessage {
1172                        id: message.id,
1173                        role: message.role,
1174                        segments: message
1175                            .segments
1176                            .iter()
1177                            .map(|segment| match segment {
1178                                MessageSegment::Text(text) => {
1179                                    SerializedMessageSegment::Text { text: text.clone() }
1180                                }
1181                                MessageSegment::Thinking { text, signature } => {
1182                                    SerializedMessageSegment::Thinking {
1183                                        text: text.clone(),
1184                                        signature: signature.clone(),
1185                                    }
1186                                }
1187                                MessageSegment::RedactedThinking(data) => {
1188                                    SerializedMessageSegment::RedactedThinking {
1189                                        data: data.clone(),
1190                                    }
1191                                }
1192                            })
1193                            .collect(),
1194                        tool_uses: this
1195                            .tool_uses_for_message(message.id, cx)
1196                            .into_iter()
1197                            .map(|tool_use| SerializedToolUse {
1198                                id: tool_use.id,
1199                                name: tool_use.name,
1200                                input: tool_use.input,
1201                            })
1202                            .collect(),
1203                        tool_results: this
1204                            .tool_results_for_message(message.id)
1205                            .into_iter()
1206                            .map(|tool_result| SerializedToolResult {
1207                                tool_use_id: tool_result.tool_use_id.clone(),
1208                                is_error: tool_result.is_error,
1209                                content: tool_result.content.clone(),
1210                                output: tool_result.output.clone(),
1211                            })
1212                            .collect(),
1213                        context: message.loaded_context.text.clone(),
1214                        creases: message
1215                            .creases
1216                            .iter()
1217                            .map(|crease| SerializedCrease {
1218                                start: crease.range.start,
1219                                end: crease.range.end,
1220                                icon_path: crease.icon_path.clone(),
1221                                label: crease.label.clone(),
1222                            })
1223                            .collect(),
1224                        is_hidden: message.is_hidden,
1225                    })
1226                    .collect(),
1227                initial_project_snapshot,
1228                cumulative_token_usage: this.cumulative_token_usage,
1229                request_token_usage: this.request_token_usage.clone(),
1230                detailed_summary_state: this.detailed_summary_rx.borrow().clone(),
1231                exceeded_window_error: this.exceeded_window_error.clone(),
1232                model: this
1233                    .configured_model
1234                    .as_ref()
1235                    .map(|model| SerializedLanguageModel {
1236                        provider: model.provider.id().0.to_string(),
1237                        model: model.model.id().0.to_string(),
1238                    }),
1239                completion_mode: Some(this.completion_mode),
1240                tool_use_limit_reached: this.tool_use_limit_reached,
1241                profile: Some(this.profile.id().clone()),
1242            })
1243        })
1244    }
1245
1246    pub fn remaining_turns(&self) -> u32 {
1247        self.remaining_turns
1248    }
1249
1250    pub fn set_remaining_turns(&mut self, remaining_turns: u32) {
1251        self.remaining_turns = remaining_turns;
1252    }
1253
1254    pub fn send_to_model(
1255        &mut self,
1256        model: Arc<dyn LanguageModel>,
1257        intent: CompletionIntent,
1258        window: Option<AnyWindowHandle>,
1259        cx: &mut Context<Self>,
1260    ) {
1261        if self.remaining_turns == 0 {
1262            return;
1263        }
1264
1265        self.remaining_turns -= 1;
1266
1267        self.flush_notifications(model.clone(), intent, cx);
1268
1269        let _checkpoint = self.finalize_pending_checkpoint(cx);
1270        self.stream_completion(
1271            self.to_completion_request(model.clone(), intent, cx),
1272            model,
1273            intent,
1274            window,
1275            cx,
1276        );
1277    }
1278
1279    pub fn retry_last_completion(
1280        &mut self,
1281        window: Option<AnyWindowHandle>,
1282        cx: &mut Context<Self>,
1283    ) {
1284        // Clear any existing error state
1285        self.retry_state = None;
1286
1287        // Use the last error context if available, otherwise fall back to configured model
1288        let (model, intent) = if let Some((model, intent)) = self.last_error_context.take() {
1289            (model, intent)
1290        } else if let Some(configured_model) = self.configured_model.as_ref() {
1291            let model = configured_model.model.clone();
1292            let intent = if self.has_pending_tool_uses() {
1293                CompletionIntent::ToolResults
1294            } else {
1295                CompletionIntent::UserPrompt
1296            };
1297            (model, intent)
1298        } else if let Some(configured_model) = self.get_or_init_configured_model(cx) {
1299            let model = configured_model.model.clone();
1300            let intent = if self.has_pending_tool_uses() {
1301                CompletionIntent::ToolResults
1302            } else {
1303                CompletionIntent::UserPrompt
1304            };
1305            (model, intent)
1306        } else {
1307            return;
1308        };
1309
1310        self.send_to_model(model, intent, window, cx);
1311    }
1312
1313    pub fn enable_burn_mode_and_retry(
1314        &mut self,
1315        window: Option<AnyWindowHandle>,
1316        cx: &mut Context<Self>,
1317    ) {
1318        self.completion_mode = CompletionMode::Burn;
1319        cx.emit(ThreadEvent::ProfileChanged);
1320        self.retry_last_completion(window, cx);
1321    }
1322
1323    pub fn used_tools_since_last_user_message(&self) -> bool {
1324        for message in self.messages.iter().rev() {
1325            if self.tool_use.message_has_tool_results(message.id) {
1326                return true;
1327            } else if message.role == Role::User {
1328                return false;
1329            }
1330        }
1331
1332        false
1333    }
1334
1335    pub fn to_completion_request(
1336        &self,
1337        model: Arc<dyn LanguageModel>,
1338        intent: CompletionIntent,
1339        cx: &mut Context<Self>,
1340    ) -> LanguageModelRequest {
1341        let mut request = LanguageModelRequest {
1342            thread_id: Some(self.id.to_string()),
1343            prompt_id: Some(self.last_prompt_id.to_string()),
1344            intent: Some(intent),
1345            mode: None,
1346            messages: vec![],
1347            tools: Vec::new(),
1348            tool_choice: None,
1349            stop: Vec::new(),
1350            temperature: AgentSettings::temperature_for_model(&model, cx),
1351            thinking_allowed: true,
1352        };
1353
1354        let available_tools = self.available_tools(cx, model.clone());
1355        let available_tool_names = available_tools
1356            .iter()
1357            .map(|tool| tool.name.clone())
1358            .collect();
1359
1360        let model_context = &ModelContext {
1361            available_tools: available_tool_names,
1362        };
1363
1364        if let Some(project_context) = self.project_context.borrow().as_ref() {
1365            match self
1366                .prompt_builder
1367                .generate_assistant_system_prompt(project_context, model_context)
1368            {
1369                Err(err) => {
1370                    let message = format!("{err:?}").into();
1371                    log::error!("{message}");
1372                    cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1373                        header: "Error generating system prompt".into(),
1374                        message,
1375                    }));
1376                }
1377                Ok(system_prompt) => {
1378                    request.messages.push(LanguageModelRequestMessage {
1379                        role: Role::System,
1380                        content: vec![MessageContent::Text(system_prompt)],
1381                        cache: true,
1382                    });
1383                }
1384            }
1385        } else {
1386            let message = "Context for system prompt unexpectedly not ready.".into();
1387            log::error!("{message}");
1388            cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1389                header: "Error generating system prompt".into(),
1390                message,
1391            }));
1392        }
1393
1394        let mut message_ix_to_cache = None;
1395        for message in &self.messages {
1396            // ui_only messages are for the UI only, not for the model
1397            if message.ui_only {
1398                continue;
1399            }
1400
1401            let mut request_message = LanguageModelRequestMessage {
1402                role: message.role,
1403                content: Vec::new(),
1404                cache: false,
1405            };
1406
1407            message
1408                .loaded_context
1409                .add_to_request_message(&mut request_message);
1410
1411            for segment in &message.segments {
1412                match segment {
1413                    MessageSegment::Text(text) => {
1414                        let text = text.trim_end();
1415                        if !text.is_empty() {
1416                            request_message
1417                                .content
1418                                .push(MessageContent::Text(text.into()));
1419                        }
1420                    }
1421                    MessageSegment::Thinking { text, signature } => {
1422                        if !text.is_empty() {
1423                            request_message.content.push(MessageContent::Thinking {
1424                                text: text.into(),
1425                                signature: signature.clone(),
1426                            });
1427                        }
1428                    }
1429                    MessageSegment::RedactedThinking(data) => {
1430                        request_message
1431                            .content
1432                            .push(MessageContent::RedactedThinking(data.clone()));
1433                    }
1434                };
1435            }
1436
1437            let mut cache_message = true;
1438            let mut tool_results_message = LanguageModelRequestMessage {
1439                role: Role::User,
1440                content: Vec::new(),
1441                cache: false,
1442            };
1443            for (tool_use, tool_result) in self.tool_use.tool_results(message.id) {
1444                if let Some(tool_result) = tool_result {
1445                    request_message
1446                        .content
1447                        .push(MessageContent::ToolUse(tool_use.clone()));
1448                    tool_results_message
1449                        .content
1450                        .push(MessageContent::ToolResult(LanguageModelToolResult {
1451                            tool_use_id: tool_use.id.clone(),
1452                            tool_name: tool_result.tool_name.clone(),
1453                            is_error: tool_result.is_error,
1454                            content: if tool_result.content.is_empty() {
1455                                // Surprisingly, the API fails if we return an empty string here.
1456                                // It thinks we are sending a tool use without a tool result.
1457                                "<Tool returned an empty string>".into()
1458                            } else {
1459                                tool_result.content.clone()
1460                            },
1461                            output: None,
1462                        }));
1463                } else {
1464                    cache_message = false;
1465                    log::debug!(
1466                        "skipped tool use {:?} because it is still pending",
1467                        tool_use
1468                    );
1469                }
1470            }
1471
1472            if cache_message {
1473                message_ix_to_cache = Some(request.messages.len());
1474            }
1475            request.messages.push(request_message);
1476
1477            if !tool_results_message.content.is_empty() {
1478                if cache_message {
1479                    message_ix_to_cache = Some(request.messages.len());
1480                }
1481                request.messages.push(tool_results_message);
1482            }
1483        }
1484
1485        // https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
1486        if let Some(message_ix_to_cache) = message_ix_to_cache {
1487            request.messages[message_ix_to_cache].cache = true;
1488        }
1489
1490        request.tools = available_tools;
1491        request.mode = if model.supports_burn_mode() {
1492            Some(self.completion_mode.into())
1493        } else {
1494            Some(CompletionMode::Normal.into())
1495        };
1496
1497        request
1498    }
1499
1500    fn to_summarize_request(
1501        &self,
1502        model: &Arc<dyn LanguageModel>,
1503        intent: CompletionIntent,
1504        added_user_message: String,
1505        cx: &App,
1506    ) -> LanguageModelRequest {
1507        let mut request = LanguageModelRequest {
1508            thread_id: None,
1509            prompt_id: None,
1510            intent: Some(intent),
1511            mode: None,
1512            messages: vec![],
1513            tools: Vec::new(),
1514            tool_choice: None,
1515            stop: Vec::new(),
1516            temperature: AgentSettings::temperature_for_model(model, cx),
1517            thinking_allowed: false,
1518        };
1519
1520        for message in &self.messages {
1521            let mut request_message = LanguageModelRequestMessage {
1522                role: message.role,
1523                content: Vec::new(),
1524                cache: false,
1525            };
1526
1527            for segment in &message.segments {
1528                match segment {
1529                    MessageSegment::Text(text) => request_message
1530                        .content
1531                        .push(MessageContent::Text(text.clone())),
1532                    MessageSegment::Thinking { .. } => {}
1533                    MessageSegment::RedactedThinking(_) => {}
1534                }
1535            }
1536
1537            if request_message.content.is_empty() {
1538                continue;
1539            }
1540
1541            request.messages.push(request_message);
1542        }
1543
1544        request.messages.push(LanguageModelRequestMessage {
1545            role: Role::User,
1546            content: vec![MessageContent::Text(added_user_message)],
1547            cache: false,
1548        });
1549
1550        request
1551    }
1552
1553    /// Insert auto-generated notifications (if any) to the thread
1554    fn flush_notifications(
1555        &mut self,
1556        model: Arc<dyn LanguageModel>,
1557        intent: CompletionIntent,
1558        cx: &mut Context<Self>,
1559    ) {
1560        match intent {
1561            CompletionIntent::UserPrompt | CompletionIntent::ToolResults => {
1562                if let Some(pending_tool_use) = self.attach_tracked_files_state(model, cx) {
1563                    cx.emit(ThreadEvent::ToolFinished {
1564                        tool_use_id: pending_tool_use.id.clone(),
1565                        pending_tool_use: Some(pending_tool_use),
1566                    });
1567                }
1568            }
1569            CompletionIntent::ThreadSummarization
1570            | CompletionIntent::ThreadContextSummarization
1571            | CompletionIntent::CreateFile
1572            | CompletionIntent::EditFile
1573            | CompletionIntent::InlineAssist
1574            | CompletionIntent::TerminalInlineAssist
1575            | CompletionIntent::GenerateGitCommitMessage => {}
1576        };
1577    }
1578
1579    fn attach_tracked_files_state(
1580        &mut self,
1581        model: Arc<dyn LanguageModel>,
1582        cx: &mut App,
1583    ) -> Option<PendingToolUse> {
1584        // Represent notification as a simulated `project_notifications` tool call
1585        let tool_name = Arc::from("project_notifications");
1586        let tool = self.tools.read(cx).tool(&tool_name, cx)?;
1587
1588        if !self.profile.is_tool_enabled(tool.source(), tool.name(), cx) {
1589            return None;
1590        }
1591
1592        if self
1593            .action_log
1594            .update(cx, |log, cx| log.unnotified_user_edits(cx).is_none())
1595        {
1596            return None;
1597        }
1598
1599        let input = serde_json::json!({});
1600        let request = Arc::new(LanguageModelRequest::default()); // unused
1601        let window = None;
1602        let tool_result = tool.run(
1603            input,
1604            request,
1605            self.project.clone(),
1606            self.action_log.clone(),
1607            model.clone(),
1608            window,
1609            cx,
1610        );
1611
1612        let tool_use_id =
1613            LanguageModelToolUseId::from(format!("project_notifications_{}", self.messages.len()));
1614
1615        let tool_use = LanguageModelToolUse {
1616            id: tool_use_id.clone(),
1617            name: tool_name.clone(),
1618            raw_input: "{}".to_string(),
1619            input: serde_json::json!({}),
1620            is_input_complete: true,
1621        };
1622
1623        let tool_output = cx.background_executor().block(tool_result.output);
1624
1625        // Attach a project_notification tool call to the latest existing
1626        // Assistant message. We cannot create a new Assistant message
1627        // because thinking models require a `thinking` block that we
1628        // cannot mock. We cannot send a notification as a normal
1629        // (non-tool-use) User message because this distracts Agent
1630        // too much.
1631        let tool_message_id = self
1632            .messages
1633            .iter()
1634            .enumerate()
1635            .rfind(|(_, message)| message.role == Role::Assistant)
1636            .map(|(_, message)| message.id)?;
1637
1638        let tool_use_metadata = ToolUseMetadata {
1639            model: model.clone(),
1640            thread_id: self.id.clone(),
1641            prompt_id: self.last_prompt_id.clone(),
1642        };
1643
1644        self.tool_use
1645            .request_tool_use(tool_message_id, tool_use, tool_use_metadata.clone(), cx);
1646
1647        let pending_tool_use = self.tool_use.insert_tool_output(
1648            tool_use_id.clone(),
1649            tool_name,
1650            tool_output,
1651            self.configured_model.as_ref(),
1652            self.completion_mode,
1653        );
1654
1655        pending_tool_use
1656    }
1657
1658    pub fn stream_completion(
1659        &mut self,
1660        request: LanguageModelRequest,
1661        model: Arc<dyn LanguageModel>,
1662        intent: CompletionIntent,
1663        window: Option<AnyWindowHandle>,
1664        cx: &mut Context<Self>,
1665    ) {
1666        self.tool_use_limit_reached = false;
1667
1668        let pending_completion_id = post_inc(&mut self.completion_count);
1669        let mut request_callback_parameters = if self.request_callback.is_some() {
1670            Some((request.clone(), Vec::new()))
1671        } else {
1672            None
1673        };
1674        let prompt_id = self.last_prompt_id.clone();
1675        let tool_use_metadata = ToolUseMetadata {
1676            model: model.clone(),
1677            thread_id: self.id.clone(),
1678            prompt_id: prompt_id.clone(),
1679        };
1680
1681        let completion_mode = request
1682            .mode
1683            .unwrap_or(cloud_llm_client::CompletionMode::Normal);
1684
1685        self.last_received_chunk_at = Some(Instant::now());
1686
1687        let task = cx.spawn(async move |thread, cx| {
1688            let stream_completion_future = model.stream_completion(request, &cx);
1689            let initial_token_usage =
1690                thread.read_with(cx, |thread, _cx| thread.cumulative_token_usage);
1691            let stream_completion = async {
1692                let mut events = stream_completion_future.await?;
1693
1694                let mut stop_reason = StopReason::EndTurn;
1695                let mut current_token_usage = TokenUsage::default();
1696
1697                thread
1698                    .update(cx, |_thread, cx| {
1699                        cx.emit(ThreadEvent::NewRequest);
1700                    })
1701                    .ok();
1702
1703                let mut request_assistant_message_id = None;
1704
1705                while let Some(event) = events.next().await {
1706                    if let Some((_, response_events)) = request_callback_parameters.as_mut() {
1707                        response_events
1708                            .push(event.as_ref().map_err(|error| error.to_string()).cloned());
1709                    }
1710
1711                    thread.update(cx, |thread, cx| {
1712                        match event? {
1713                            LanguageModelCompletionEvent::StartMessage { .. } => {
1714                                request_assistant_message_id =
1715                                    Some(thread.insert_assistant_message(
1716                                        vec![MessageSegment::Text(String::new())],
1717                                        cx,
1718                                    ));
1719                            }
1720                            LanguageModelCompletionEvent::Stop(reason) => {
1721                                stop_reason = reason;
1722                            }
1723                            LanguageModelCompletionEvent::UsageUpdate(token_usage) => {
1724                                thread.update_token_usage_at_last_message(token_usage);
1725                                thread.cumulative_token_usage = thread.cumulative_token_usage
1726                                    + token_usage
1727                                    - current_token_usage;
1728                                current_token_usage = token_usage;
1729                            }
1730                            LanguageModelCompletionEvent::Text(chunk) => {
1731                                thread.received_chunk();
1732
1733                                cx.emit(ThreadEvent::ReceivedTextChunk);
1734                                if let Some(last_message) = thread.messages.last_mut() {
1735                                    if last_message.role == Role::Assistant
1736                                        && !thread.tool_use.has_tool_results(last_message.id)
1737                                    {
1738                                        last_message.push_text(&chunk);
1739                                        cx.emit(ThreadEvent::StreamedAssistantText(
1740                                            last_message.id,
1741                                            chunk,
1742                                        ));
1743                                    } else {
1744                                        // If we won't have an Assistant message yet, assume this chunk marks the beginning
1745                                        // of a new Assistant response.
1746                                        //
1747                                        // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1748                                        // will result in duplicating the text of the chunk in the rendered Markdown.
1749                                        request_assistant_message_id =
1750                                            Some(thread.insert_assistant_message(
1751                                                vec![MessageSegment::Text(chunk.to_string())],
1752                                                cx,
1753                                            ));
1754                                    };
1755                                }
1756                            }
1757                            LanguageModelCompletionEvent::Thinking {
1758                                text: chunk,
1759                                signature,
1760                            } => {
1761                                thread.received_chunk();
1762
1763                                if let Some(last_message) = thread.messages.last_mut() {
1764                                    if last_message.role == Role::Assistant
1765                                        && !thread.tool_use.has_tool_results(last_message.id)
1766                                    {
1767                                        last_message.push_thinking(&chunk, signature);
1768                                        cx.emit(ThreadEvent::StreamedAssistantThinking(
1769                                            last_message.id,
1770                                            chunk,
1771                                        ));
1772                                    } else {
1773                                        // If we won't have an Assistant message yet, assume this chunk marks the beginning
1774                                        // of a new Assistant response.
1775                                        //
1776                                        // Importantly: We do *not* want to emit a `StreamedAssistantText` event here, as it
1777                                        // will result in duplicating the text of the chunk in the rendered Markdown.
1778                                        request_assistant_message_id =
1779                                            Some(thread.insert_assistant_message(
1780                                                vec![MessageSegment::Thinking {
1781                                                    text: chunk.to_string(),
1782                                                    signature,
1783                                                }],
1784                                                cx,
1785                                            ));
1786                                    };
1787                                }
1788                            }
1789                            LanguageModelCompletionEvent::RedactedThinking { data } => {
1790                                thread.received_chunk();
1791
1792                                if let Some(last_message) = thread.messages.last_mut() {
1793                                    if last_message.role == Role::Assistant
1794                                        && !thread.tool_use.has_tool_results(last_message.id)
1795                                    {
1796                                        last_message.push_redacted_thinking(data);
1797                                    } else {
1798                                        request_assistant_message_id =
1799                                            Some(thread.insert_assistant_message(
1800                                                vec![MessageSegment::RedactedThinking(data)],
1801                                                cx,
1802                                            ));
1803                                    };
1804                                }
1805                            }
1806                            LanguageModelCompletionEvent::ToolUse(tool_use) => {
1807                                let last_assistant_message_id = request_assistant_message_id
1808                                    .unwrap_or_else(|| {
1809                                        let new_assistant_message_id =
1810                                            thread.insert_assistant_message(vec![], cx);
1811                                        request_assistant_message_id =
1812                                            Some(new_assistant_message_id);
1813                                        new_assistant_message_id
1814                                    });
1815
1816                                let tool_use_id = tool_use.id.clone();
1817                                let streamed_input = if tool_use.is_input_complete {
1818                                    None
1819                                } else {
1820                                    Some((&tool_use.input).clone())
1821                                };
1822
1823                                let ui_text = thread.tool_use.request_tool_use(
1824                                    last_assistant_message_id,
1825                                    tool_use,
1826                                    tool_use_metadata.clone(),
1827                                    cx,
1828                                );
1829
1830                                if let Some(input) = streamed_input {
1831                                    cx.emit(ThreadEvent::StreamedToolUse {
1832                                        tool_use_id,
1833                                        ui_text,
1834                                        input,
1835                                    });
1836                                }
1837                            }
1838                            LanguageModelCompletionEvent::ToolUseJsonParseError {
1839                                id,
1840                                tool_name,
1841                                raw_input: invalid_input_json,
1842                                json_parse_error,
1843                            } => {
1844                                thread.receive_invalid_tool_json(
1845                                    id,
1846                                    tool_name,
1847                                    invalid_input_json,
1848                                    json_parse_error,
1849                                    window,
1850                                    cx,
1851                                );
1852                            }
1853                            LanguageModelCompletionEvent::StatusUpdate(status_update) => {
1854                                if let Some(completion) = thread
1855                                    .pending_completions
1856                                    .iter_mut()
1857                                    .find(|completion| completion.id == pending_completion_id)
1858                                {
1859                                    match status_update {
1860                                        CompletionRequestStatus::Queued { position } => {
1861                                            completion.queue_state =
1862                                                QueueState::Queued { position };
1863                                        }
1864                                        CompletionRequestStatus::Started => {
1865                                            completion.queue_state = QueueState::Started;
1866                                        }
1867                                        CompletionRequestStatus::Failed {
1868                                            code,
1869                                            message,
1870                                            request_id: _,
1871                                            retry_after,
1872                                        } => {
1873                                            return Err(
1874                                                LanguageModelCompletionError::from_cloud_failure(
1875                                                    model.upstream_provider_name(),
1876                                                    code,
1877                                                    message,
1878                                                    retry_after.map(Duration::from_secs_f64),
1879                                                ),
1880                                            );
1881                                        }
1882                                        CompletionRequestStatus::UsageUpdated { amount, limit } => {
1883                                            thread.update_model_request_usage(
1884                                                amount as u32,
1885                                                limit,
1886                                                cx,
1887                                            );
1888                                        }
1889                                        CompletionRequestStatus::ToolUseLimitReached => {
1890                                            thread.tool_use_limit_reached = true;
1891                                            cx.emit(ThreadEvent::ToolUseLimitReached);
1892                                        }
1893                                    }
1894                                }
1895                            }
1896                        }
1897
1898                        thread.touch_updated_at();
1899                        cx.emit(ThreadEvent::StreamedCompletion);
1900                        cx.notify();
1901
1902                        thread.auto_capture_telemetry(cx);
1903                        Ok(())
1904                    })??;
1905
1906                    smol::future::yield_now().await;
1907                }
1908
1909                thread.update(cx, |thread, cx| {
1910                    thread.last_received_chunk_at = None;
1911                    thread
1912                        .pending_completions
1913                        .retain(|completion| completion.id != pending_completion_id);
1914
1915                    // If there is a response without tool use, summarize the message. Otherwise,
1916                    // allow two tool uses before summarizing.
1917                    if matches!(thread.summary, ThreadSummary::Pending)
1918                        && thread.messages.len() >= 2
1919                        && (!thread.has_pending_tool_uses() || thread.messages.len() >= 6)
1920                    {
1921                        thread.summarize(cx);
1922                    }
1923                })?;
1924
1925                anyhow::Ok(stop_reason)
1926            };
1927
1928            let result = stream_completion.await;
1929            let mut retry_scheduled = false;
1930
1931            thread
1932                .update(cx, |thread, cx| {
1933                    thread.finalize_pending_checkpoint(cx);
1934                    match result.as_ref() {
1935                        Ok(stop_reason) => {
1936                            match stop_reason {
1937                                StopReason::ToolUse => {
1938                                    let tool_uses =
1939                                        thread.use_pending_tools(window, model.clone(), cx);
1940                                    cx.emit(ThreadEvent::UsePendingTools { tool_uses });
1941                                }
1942                                StopReason::EndTurn | StopReason::MaxTokens => {
1943                                    thread.project.update(cx, |project, cx| {
1944                                        project.set_agent_location(None, cx);
1945                                    });
1946                                }
1947                                StopReason::Refusal => {
1948                                    thread.project.update(cx, |project, cx| {
1949                                        project.set_agent_location(None, cx);
1950                                    });
1951
1952                                    // Remove the turn that was refused.
1953                                    //
1954                                    // https://docs.anthropic.com/en/docs/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals#reset-context-after-refusal
1955                                    {
1956                                        let mut messages_to_remove = Vec::new();
1957
1958                                        for (ix, message) in
1959                                            thread.messages.iter().enumerate().rev()
1960                                        {
1961                                            messages_to_remove.push(message.id);
1962
1963                                            if message.role == Role::User {
1964                                                if ix == 0 {
1965                                                    break;
1966                                                }
1967
1968                                                if let Some(prev_message) =
1969                                                    thread.messages.get(ix - 1)
1970                                                {
1971                                                    if prev_message.role == Role::Assistant {
1972                                                        break;
1973                                                    }
1974                                                }
1975                                            }
1976                                        }
1977
1978                                        for message_id in messages_to_remove {
1979                                            thread.delete_message(message_id, cx);
1980                                        }
1981                                    }
1982
1983                                    cx.emit(ThreadEvent::ShowError(ThreadError::Message {
1984                                        header: "Language model refusal".into(),
1985                                        message:
1986                                            "Model refused to generate content for safety reasons."
1987                                                .into(),
1988                                    }));
1989                                }
1990                            }
1991
1992                            // We successfully completed, so cancel any remaining retries.
1993                            thread.retry_state = None;
1994                        }
1995                        Err(error) => {
1996                            thread.project.update(cx, |project, cx| {
1997                                project.set_agent_location(None, cx);
1998                            });
1999
2000                            if error.is::<PaymentRequiredError>() {
2001                                cx.emit(ThreadEvent::ShowError(ThreadError::PaymentRequired));
2002                            } else if let Some(error) =
2003                                error.downcast_ref::<ModelRequestLimitReachedError>()
2004                            {
2005                                cx.emit(ThreadEvent::ShowError(
2006                                    ThreadError::ModelRequestLimitReached { plan: error.plan },
2007                                ));
2008                            } else if let Some(completion_error) =
2009                                error.downcast_ref::<LanguageModelCompletionError>()
2010                            {
2011                                match &completion_error {
2012                                    LanguageModelCompletionError::PromptTooLarge {
2013                                        tokens, ..
2014                                    } => {
2015                                        let tokens = tokens.unwrap_or_else(|| {
2016                                            // We didn't get an exact token count from the API, so fall back on our estimate.
2017                                            thread
2018                                                .total_token_usage()
2019                                                .map(|usage| usage.total)
2020                                                .unwrap_or(0)
2021                                                // We know the context window was exceeded in practice, so if our estimate was
2022                                                // lower than max tokens, the estimate was wrong; return that we exceeded by 1.
2023                                                .max(
2024                                                    model
2025                                                        .max_token_count_for_mode(completion_mode)
2026                                                        .saturating_add(1),
2027                                                )
2028                                        });
2029                                        thread.exceeded_window_error = Some(ExceededWindowError {
2030                                            model_id: model.id(),
2031                                            token_count: tokens,
2032                                        });
2033                                        cx.notify();
2034                                    }
2035                                    _ => {
2036                                        if let Some(retry_strategy) =
2037                                            Thread::get_retry_strategy(completion_error)
2038                                        {
2039                                            log::info!(
2040                                                "Retrying with {:?} for language model completion error {:?}",
2041                                                retry_strategy,
2042                                                completion_error
2043                                            );
2044
2045                                            retry_scheduled = thread
2046                                                .handle_retryable_error_with_delay(
2047                                                    &completion_error,
2048                                                    Some(retry_strategy),
2049                                                    model.clone(),
2050                                                    intent,
2051                                                    window,
2052                                                    cx,
2053                                                );
2054                                        }
2055                                    }
2056                                }
2057                            }
2058
2059                            if !retry_scheduled {
2060                                thread.cancel_last_completion(window, cx);
2061                            }
2062                        }
2063                    }
2064
2065                    if !retry_scheduled {
2066                        cx.emit(ThreadEvent::Stopped(result.map_err(Arc::new)));
2067                    }
2068
2069                    if let Some((request_callback, (request, response_events))) = thread
2070                        .request_callback
2071                        .as_mut()
2072                        .zip(request_callback_parameters.as_ref())
2073                    {
2074                        request_callback(request, response_events);
2075                    }
2076
2077                    thread.auto_capture_telemetry(cx);
2078
2079                    if let Ok(initial_usage) = initial_token_usage {
2080                        let usage = thread.cumulative_token_usage - initial_usage;
2081
2082                        telemetry::event!(
2083                            "Assistant Thread Completion",
2084                            thread_id = thread.id().to_string(),
2085                            prompt_id = prompt_id,
2086                            model = model.telemetry_id(),
2087                            model_provider = model.provider_id().to_string(),
2088                            input_tokens = usage.input_tokens,
2089                            output_tokens = usage.output_tokens,
2090                            cache_creation_input_tokens = usage.cache_creation_input_tokens,
2091                            cache_read_input_tokens = usage.cache_read_input_tokens,
2092                        );
2093                    }
2094                })
2095                .ok();
2096        });
2097
2098        self.pending_completions.push(PendingCompletion {
2099            id: pending_completion_id,
2100            queue_state: QueueState::Sending,
2101            _task: task,
2102        });
2103    }
2104
2105    pub fn summarize(&mut self, cx: &mut Context<Self>) {
2106        let Some(model) = LanguageModelRegistry::read_global(cx).thread_summary_model() else {
2107            println!("No thread summary model");
2108            return;
2109        };
2110
2111        if !model.provider.is_authenticated(cx) {
2112            return;
2113        }
2114
2115        let added_user_message = include_str!("./prompts/summarize_thread_prompt.txt");
2116
2117        let request = self.to_summarize_request(
2118            &model.model,
2119            CompletionIntent::ThreadSummarization,
2120            added_user_message.into(),
2121            cx,
2122        );
2123
2124        self.summary = ThreadSummary::Generating;
2125
2126        self.pending_summary = cx.spawn(async move |this, cx| {
2127            let result = async {
2128                let mut messages = model.model.stream_completion(request, &cx).await?;
2129
2130                let mut new_summary = String::new();
2131                while let Some(event) = messages.next().await {
2132                    let Ok(event) = event else {
2133                        continue;
2134                    };
2135                    let text = match event {
2136                        LanguageModelCompletionEvent::Text(text) => text,
2137                        LanguageModelCompletionEvent::StatusUpdate(
2138                            CompletionRequestStatus::UsageUpdated { amount, limit },
2139                        ) => {
2140                            this.update(cx, |thread, cx| {
2141                                thread.update_model_request_usage(amount as u32, limit, cx);
2142                            })?;
2143                            continue;
2144                        }
2145                        _ => continue,
2146                    };
2147
2148                    let mut lines = text.lines();
2149                    new_summary.extend(lines.next());
2150
2151                    // Stop if the LLM generated multiple lines.
2152                    if lines.next().is_some() {
2153                        break;
2154                    }
2155                }
2156
2157                anyhow::Ok(new_summary)
2158            }
2159            .await;
2160
2161            this.update(cx, |this, cx| {
2162                match result {
2163                    Ok(new_summary) => {
2164                        if new_summary.is_empty() {
2165                            this.summary = ThreadSummary::Error;
2166                        } else {
2167                            this.summary = ThreadSummary::Ready(new_summary.into());
2168                        }
2169                    }
2170                    Err(err) => {
2171                        this.summary = ThreadSummary::Error;
2172                        log::error!("Failed to generate thread summary: {}", err);
2173                    }
2174                }
2175                cx.emit(ThreadEvent::SummaryGenerated);
2176            })
2177            .log_err()?;
2178
2179            Some(())
2180        });
2181    }
2182
2183    fn get_retry_strategy(error: &LanguageModelCompletionError) -> Option<RetryStrategy> {
2184        use LanguageModelCompletionError::*;
2185
2186        // General strategy here:
2187        // - If retrying won't help (e.g. invalid API key or payload too large), return None so we don't retry at all.
2188        // - If it's a time-based issue (e.g. server overloaded, rate limit exceeded), retry up to 4 times with exponential backoff.
2189        // - If it's an issue that *might* be fixed by retrying (e.g. internal server error), retry up to 3 times.
2190        match error {
2191            HttpResponseError {
2192                status_code: StatusCode::TOO_MANY_REQUESTS,
2193                ..
2194            } => Some(RetryStrategy::ExponentialBackoff {
2195                initial_delay: BASE_RETRY_DELAY,
2196                max_attempts: MAX_RETRY_ATTEMPTS,
2197            }),
2198            ServerOverloaded { retry_after, .. } | RateLimitExceeded { retry_after, .. } => {
2199                Some(RetryStrategy::Fixed {
2200                    delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2201                    max_attempts: MAX_RETRY_ATTEMPTS,
2202                })
2203            }
2204            UpstreamProviderError {
2205                status,
2206                retry_after,
2207                ..
2208            } => match *status {
2209                StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => {
2210                    Some(RetryStrategy::Fixed {
2211                        delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2212                        max_attempts: MAX_RETRY_ATTEMPTS,
2213                    })
2214                }
2215                StatusCode::INTERNAL_SERVER_ERROR => Some(RetryStrategy::Fixed {
2216                    delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2217                    // Internal Server Error could be anything, retry up to 3 times.
2218                    max_attempts: 3,
2219                }),
2220                status => {
2221                    // There is no StatusCode variant for the unofficial HTTP 529 ("The service is overloaded"),
2222                    // but we frequently get them in practice. See https://http.dev/529
2223                    if status.as_u16() == 529 {
2224                        Some(RetryStrategy::Fixed {
2225                            delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2226                            max_attempts: MAX_RETRY_ATTEMPTS,
2227                        })
2228                    } else {
2229                        Some(RetryStrategy::Fixed {
2230                            delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2231                            max_attempts: 2,
2232                        })
2233                    }
2234                }
2235            },
2236            ApiInternalServerError { .. } => Some(RetryStrategy::Fixed {
2237                delay: BASE_RETRY_DELAY,
2238                max_attempts: 3,
2239            }),
2240            ApiReadResponseError { .. }
2241            | HttpSend { .. }
2242            | DeserializeResponse { .. }
2243            | BadRequestFormat { .. } => Some(RetryStrategy::Fixed {
2244                delay: BASE_RETRY_DELAY,
2245                max_attempts: 3,
2246            }),
2247            // Retrying these errors definitely shouldn't help.
2248            HttpResponseError {
2249                status_code:
2250                    StatusCode::PAYLOAD_TOO_LARGE | StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED,
2251                ..
2252            }
2253            | AuthenticationError { .. }
2254            | PermissionError { .. }
2255            | NoApiKey { .. }
2256            | ApiEndpointNotFound { .. }
2257            | PromptTooLarge { .. } => None,
2258            // These errors might be transient, so retry them
2259            SerializeRequest { .. } | BuildRequestBody { .. } => Some(RetryStrategy::Fixed {
2260                delay: BASE_RETRY_DELAY,
2261                max_attempts: 1,
2262            }),
2263            // Retry all other 4xx and 5xx errors once.
2264            HttpResponseError { status_code, .. }
2265                if status_code.is_client_error() || status_code.is_server_error() =>
2266            {
2267                Some(RetryStrategy::Fixed {
2268                    delay: BASE_RETRY_DELAY,
2269                    max_attempts: 3,
2270                })
2271            }
2272            // Conservatively assume that any other errors are non-retryable
2273            HttpResponseError { .. } | Other(..) => Some(RetryStrategy::Fixed {
2274                delay: BASE_RETRY_DELAY,
2275                max_attempts: 2,
2276            }),
2277        }
2278    }
2279
2280    fn handle_retryable_error_with_delay(
2281        &mut self,
2282        error: &LanguageModelCompletionError,
2283        strategy: Option<RetryStrategy>,
2284        model: Arc<dyn LanguageModel>,
2285        intent: CompletionIntent,
2286        window: Option<AnyWindowHandle>,
2287        cx: &mut Context<Self>,
2288    ) -> bool {
2289        // Store context for the Retry button
2290        self.last_error_context = Some((model.clone(), intent));
2291
2292        // Only auto-retry if Burn Mode is enabled
2293        if self.completion_mode != CompletionMode::Burn {
2294            // Show error with retry options
2295            cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError {
2296                message: format!(
2297                    "{}\n\nTo automatically retry when similar errors happen, enable Burn Mode.",
2298                    error
2299                )
2300                .into(),
2301                can_enable_burn_mode: true,
2302            }));
2303            return false;
2304        }
2305
2306        let Some(strategy) = strategy.or_else(|| Self::get_retry_strategy(error)) else {
2307            return false;
2308        };
2309
2310        let max_attempts = match &strategy {
2311            RetryStrategy::ExponentialBackoff { max_attempts, .. } => *max_attempts,
2312            RetryStrategy::Fixed { max_attempts, .. } => *max_attempts,
2313        };
2314
2315        let retry_state = self.retry_state.get_or_insert(RetryState {
2316            attempt: 0,
2317            max_attempts,
2318            intent,
2319        });
2320
2321        retry_state.attempt += 1;
2322        let attempt = retry_state.attempt;
2323        let max_attempts = retry_state.max_attempts;
2324        let intent = retry_state.intent;
2325
2326        if attempt <= max_attempts {
2327            let delay = match &strategy {
2328                RetryStrategy::ExponentialBackoff { initial_delay, .. } => {
2329                    let delay_secs = initial_delay.as_secs() * 2u64.pow((attempt - 1) as u32);
2330                    Duration::from_secs(delay_secs)
2331                }
2332                RetryStrategy::Fixed { delay, .. } => *delay,
2333            };
2334
2335            // Add a transient message to inform the user
2336            let delay_secs = delay.as_secs();
2337            let retry_message = if max_attempts == 1 {
2338                format!("{error}. Retrying in {delay_secs} seconds...")
2339            } else {
2340                format!(
2341                    "{error}. Retrying (attempt {attempt} of {max_attempts}) \
2342                    in {delay_secs} seconds..."
2343                )
2344            };
2345            log::warn!(
2346                "Retrying completion request (attempt {attempt} of {max_attempts}) \
2347                in {delay_secs} seconds: {error:?}",
2348            );
2349
2350            // Add a UI-only message instead of a regular message
2351            let id = self.next_message_id.post_inc();
2352            self.messages.push(Message {
2353                id,
2354                role: Role::System,
2355                segments: vec![MessageSegment::Text(retry_message)],
2356                loaded_context: LoadedContext::default(),
2357                creases: Vec::new(),
2358                is_hidden: false,
2359                ui_only: true,
2360            });
2361            cx.emit(ThreadEvent::MessageAdded(id));
2362
2363            // Schedule the retry
2364            let thread_handle = cx.entity().downgrade();
2365
2366            cx.spawn(async move |_thread, cx| {
2367                cx.background_executor().timer(delay).await;
2368
2369                thread_handle
2370                    .update(cx, |thread, cx| {
2371                        // Retry the completion
2372                        thread.send_to_model(model, intent, window, cx);
2373                    })
2374                    .log_err();
2375            })
2376            .detach();
2377
2378            true
2379        } else {
2380            // Max retries exceeded
2381            self.retry_state = None;
2382
2383            // Stop generating since we're giving up on retrying.
2384            self.pending_completions.clear();
2385
2386            // Show error alongside a Retry button, but no
2387            // Enable Burn Mode button (since it's already enabled)
2388            cx.emit(ThreadEvent::ShowError(ThreadError::RetryableError {
2389                message: format!("Failed after retrying: {}", error).into(),
2390                can_enable_burn_mode: false,
2391            }));
2392
2393            false
2394        }
2395    }
2396
2397    pub fn start_generating_detailed_summary_if_needed(
2398        &mut self,
2399        thread_store: WeakEntity<ThreadStore>,
2400        cx: &mut Context<Self>,
2401    ) {
2402        let Some(last_message_id) = self.messages.last().map(|message| message.id) else {
2403            return;
2404        };
2405
2406        match &*self.detailed_summary_rx.borrow() {
2407            DetailedSummaryState::Generating { message_id, .. }
2408            | DetailedSummaryState::Generated { message_id, .. }
2409                if *message_id == last_message_id =>
2410            {
2411                // Already up-to-date
2412                return;
2413            }
2414            _ => {}
2415        }
2416
2417        let Some(ConfiguredModel { model, provider }) =
2418            LanguageModelRegistry::read_global(cx).thread_summary_model()
2419        else {
2420            return;
2421        };
2422
2423        if !provider.is_authenticated(cx) {
2424            return;
2425        }
2426
2427        let added_user_message = include_str!("./prompts/summarize_thread_detailed_prompt.txt");
2428
2429        let request = self.to_summarize_request(
2430            &model,
2431            CompletionIntent::ThreadContextSummarization,
2432            added_user_message.into(),
2433            cx,
2434        );
2435
2436        *self.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generating {
2437            message_id: last_message_id,
2438        };
2439
2440        // Replace the detailed summarization task if there is one, cancelling it. It would probably
2441        // be better to allow the old task to complete, but this would require logic for choosing
2442        // which result to prefer (the old task could complete after the new one, resulting in a
2443        // stale summary).
2444        self.detailed_summary_task = cx.spawn(async move |thread, cx| {
2445            let stream = model.stream_completion_text(request, &cx);
2446            let Some(mut messages) = stream.await.log_err() else {
2447                thread
2448                    .update(cx, |thread, _cx| {
2449                        *thread.detailed_summary_tx.borrow_mut() =
2450                            DetailedSummaryState::NotGenerated;
2451                    })
2452                    .ok()?;
2453                return None;
2454            };
2455
2456            let mut new_detailed_summary = String::new();
2457
2458            while let Some(chunk) = messages.stream.next().await {
2459                if let Some(chunk) = chunk.log_err() {
2460                    new_detailed_summary.push_str(&chunk);
2461                }
2462            }
2463
2464            thread
2465                .update(cx, |thread, _cx| {
2466                    *thread.detailed_summary_tx.borrow_mut() = DetailedSummaryState::Generated {
2467                        text: new_detailed_summary.into(),
2468                        message_id: last_message_id,
2469                    };
2470                })
2471                .ok()?;
2472
2473            // Save thread so its summary can be reused later
2474            if let Some(thread) = thread.upgrade() {
2475                if let Ok(Ok(save_task)) = cx.update(|cx| {
2476                    thread_store
2477                        .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
2478                }) {
2479                    save_task.await.log_err();
2480                }
2481            }
2482
2483            Some(())
2484        });
2485    }
2486
2487    pub async fn wait_for_detailed_summary_or_text(
2488        this: &Entity<Self>,
2489        cx: &mut AsyncApp,
2490    ) -> Option<SharedString> {
2491        let mut detailed_summary_rx = this
2492            .read_with(cx, |this, _cx| this.detailed_summary_rx.clone())
2493            .ok()?;
2494        loop {
2495            match detailed_summary_rx.recv().await? {
2496                DetailedSummaryState::Generating { .. } => {}
2497                DetailedSummaryState::NotGenerated => {
2498                    return this.read_with(cx, |this, _cx| this.text().into()).ok();
2499                }
2500                DetailedSummaryState::Generated { text, .. } => return Some(text),
2501            }
2502        }
2503    }
2504
2505    pub fn latest_detailed_summary_or_text(&self) -> SharedString {
2506        self.detailed_summary_rx
2507            .borrow()
2508            .text()
2509            .unwrap_or_else(|| self.text().into())
2510    }
2511
2512    pub fn is_generating_detailed_summary(&self) -> bool {
2513        matches!(
2514            &*self.detailed_summary_rx.borrow(),
2515            DetailedSummaryState::Generating { .. }
2516        )
2517    }
2518
2519    pub fn use_pending_tools(
2520        &mut self,
2521        window: Option<AnyWindowHandle>,
2522        model: Arc<dyn LanguageModel>,
2523        cx: &mut Context<Self>,
2524    ) -> Vec<PendingToolUse> {
2525        self.auto_capture_telemetry(cx);
2526        let request =
2527            Arc::new(self.to_completion_request(model.clone(), CompletionIntent::ToolResults, cx));
2528        let pending_tool_uses = self
2529            .tool_use
2530            .pending_tool_uses()
2531            .into_iter()
2532            .filter(|tool_use| tool_use.status.is_idle())
2533            .cloned()
2534            .collect::<Vec<_>>();
2535
2536        for tool_use in pending_tool_uses.iter() {
2537            self.use_pending_tool(tool_use.clone(), request.clone(), model.clone(), window, cx);
2538        }
2539
2540        pending_tool_uses
2541    }
2542
2543    fn use_pending_tool(
2544        &mut self,
2545        tool_use: PendingToolUse,
2546        request: Arc<LanguageModelRequest>,
2547        model: Arc<dyn LanguageModel>,
2548        window: Option<AnyWindowHandle>,
2549        cx: &mut Context<Self>,
2550    ) {
2551        let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) else {
2552            return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
2553        };
2554
2555        if !self.profile.is_tool_enabled(tool.source(), tool.name(), cx) {
2556            return self.handle_hallucinated_tool_use(tool_use.id, tool_use.name, window, cx);
2557        }
2558
2559        if tool.needs_confirmation(&tool_use.input, &self.project, cx)
2560            && !AgentSettings::get_global(cx).always_allow_tool_actions
2561        {
2562            self.tool_use.confirm_tool_use(
2563                tool_use.id,
2564                tool_use.ui_text,
2565                tool_use.input,
2566                request,
2567                tool,
2568            );
2569            cx.emit(ThreadEvent::ToolConfirmationNeeded);
2570        } else {
2571            self.run_tool(
2572                tool_use.id,
2573                tool_use.ui_text,
2574                tool_use.input,
2575                request,
2576                tool,
2577                model,
2578                window,
2579                cx,
2580            );
2581        }
2582    }
2583
2584    pub fn handle_hallucinated_tool_use(
2585        &mut self,
2586        tool_use_id: LanguageModelToolUseId,
2587        hallucinated_tool_name: Arc<str>,
2588        window: Option<AnyWindowHandle>,
2589        cx: &mut Context<Thread>,
2590    ) {
2591        let available_tools = self.profile.enabled_tools(cx);
2592
2593        let tool_list = available_tools
2594            .iter()
2595            .map(|(name, tool)| format!("- {}: {}", name, tool.description()))
2596            .collect::<Vec<_>>()
2597            .join("\n");
2598
2599        let error_message = format!(
2600            "The tool '{}' doesn't exist or is not enabled. Available tools:\n{}",
2601            hallucinated_tool_name, tool_list
2602        );
2603
2604        let pending_tool_use = self.tool_use.insert_tool_output(
2605            tool_use_id.clone(),
2606            hallucinated_tool_name,
2607            Err(anyhow!("Missing tool call: {error_message}")),
2608            self.configured_model.as_ref(),
2609            self.completion_mode,
2610        );
2611
2612        cx.emit(ThreadEvent::MissingToolUse {
2613            tool_use_id: tool_use_id.clone(),
2614            ui_text: error_message.into(),
2615        });
2616
2617        self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2618    }
2619
2620    pub fn receive_invalid_tool_json(
2621        &mut self,
2622        tool_use_id: LanguageModelToolUseId,
2623        tool_name: Arc<str>,
2624        invalid_json: Arc<str>,
2625        error: String,
2626        window: Option<AnyWindowHandle>,
2627        cx: &mut Context<Thread>,
2628    ) {
2629        log::error!("The model returned invalid input JSON: {invalid_json}");
2630
2631        let pending_tool_use = self.tool_use.insert_tool_output(
2632            tool_use_id.clone(),
2633            tool_name,
2634            Err(anyhow!("Error parsing input JSON: {error}")),
2635            self.configured_model.as_ref(),
2636            self.completion_mode,
2637        );
2638        let ui_text = if let Some(pending_tool_use) = &pending_tool_use {
2639            pending_tool_use.ui_text.clone()
2640        } else {
2641            log::error!(
2642                "There was no pending tool use for tool use {tool_use_id}, even though it finished (with invalid input JSON)."
2643            );
2644            format!("Unknown tool {}", tool_use_id).into()
2645        };
2646
2647        cx.emit(ThreadEvent::InvalidToolInput {
2648            tool_use_id: tool_use_id.clone(),
2649            ui_text,
2650            invalid_input_json: invalid_json,
2651        });
2652
2653        self.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2654    }
2655
2656    pub fn run_tool(
2657        &mut self,
2658        tool_use_id: LanguageModelToolUseId,
2659        ui_text: impl Into<SharedString>,
2660        input: serde_json::Value,
2661        request: Arc<LanguageModelRequest>,
2662        tool: Arc<dyn Tool>,
2663        model: Arc<dyn LanguageModel>,
2664        window: Option<AnyWindowHandle>,
2665        cx: &mut Context<Thread>,
2666    ) {
2667        let task =
2668            self.spawn_tool_use(tool_use_id.clone(), request, input, tool, model, window, cx);
2669        self.tool_use
2670            .run_pending_tool(tool_use_id, ui_text.into(), task);
2671    }
2672
2673    fn spawn_tool_use(
2674        &mut self,
2675        tool_use_id: LanguageModelToolUseId,
2676        request: Arc<LanguageModelRequest>,
2677        input: serde_json::Value,
2678        tool: Arc<dyn Tool>,
2679        model: Arc<dyn LanguageModel>,
2680        window: Option<AnyWindowHandle>,
2681        cx: &mut Context<Thread>,
2682    ) -> Task<()> {
2683        let tool_name: Arc<str> = tool.name().into();
2684
2685        let tool_result = tool.run(
2686            input,
2687            request,
2688            self.project.clone(),
2689            self.action_log.clone(),
2690            model,
2691            window,
2692            cx,
2693        );
2694
2695        // Store the card separately if it exists
2696        if let Some(card) = tool_result.card.clone() {
2697            self.tool_use
2698                .insert_tool_result_card(tool_use_id.clone(), card);
2699        }
2700
2701        cx.spawn({
2702            async move |thread: WeakEntity<Thread>, cx| {
2703                let output = tool_result.output.await;
2704
2705                thread
2706                    .update(cx, |thread, cx| {
2707                        let pending_tool_use = thread.tool_use.insert_tool_output(
2708                            tool_use_id.clone(),
2709                            tool_name,
2710                            output,
2711                            thread.configured_model.as_ref(),
2712                            thread.completion_mode,
2713                        );
2714                        thread.tool_finished(tool_use_id, pending_tool_use, false, window, cx);
2715                    })
2716                    .ok();
2717            }
2718        })
2719    }
2720
2721    fn tool_finished(
2722        &mut self,
2723        tool_use_id: LanguageModelToolUseId,
2724        pending_tool_use: Option<PendingToolUse>,
2725        canceled: bool,
2726        window: Option<AnyWindowHandle>,
2727        cx: &mut Context<Self>,
2728    ) {
2729        if self.all_tools_finished() {
2730            if let Some(ConfiguredModel { model, .. }) = self.configured_model.as_ref() {
2731                if !canceled {
2732                    self.send_to_model(model.clone(), CompletionIntent::ToolResults, window, cx);
2733                }
2734                self.auto_capture_telemetry(cx);
2735            }
2736        }
2737
2738        cx.emit(ThreadEvent::ToolFinished {
2739            tool_use_id,
2740            pending_tool_use,
2741        });
2742    }
2743
2744    /// Cancels the last pending completion, if there are any pending.
2745    ///
2746    /// Returns whether a completion was canceled.
2747    pub fn cancel_last_completion(
2748        &mut self,
2749        window: Option<AnyWindowHandle>,
2750        cx: &mut Context<Self>,
2751    ) -> bool {
2752        let mut canceled = self.pending_completions.pop().is_some() || self.retry_state.is_some();
2753
2754        self.retry_state = None;
2755
2756        for pending_tool_use in self.tool_use.cancel_pending() {
2757            canceled = true;
2758            self.tool_finished(
2759                pending_tool_use.id.clone(),
2760                Some(pending_tool_use),
2761                true,
2762                window,
2763                cx,
2764            );
2765        }
2766
2767        if canceled {
2768            cx.emit(ThreadEvent::CompletionCanceled);
2769
2770            // When canceled, we always want to insert the checkpoint.
2771            // (We skip over finalize_pending_checkpoint, because it
2772            // would conclude we didn't have anything to insert here.)
2773            if let Some(checkpoint) = self.pending_checkpoint.take() {
2774                self.insert_checkpoint(checkpoint, cx);
2775            }
2776        } else {
2777            self.finalize_pending_checkpoint(cx);
2778        }
2779
2780        canceled
2781    }
2782
2783    /// Signals that any in-progress editing should be canceled.
2784    ///
2785    /// This method is used to notify listeners (like ActiveThread) that
2786    /// they should cancel any editing operations.
2787    pub fn cancel_editing(&mut self, cx: &mut Context<Self>) {
2788        cx.emit(ThreadEvent::CancelEditing);
2789    }
2790
2791    pub fn feedback(&self) -> Option<ThreadFeedback> {
2792        self.feedback
2793    }
2794
2795    pub fn message_feedback(&self, message_id: MessageId) -> Option<ThreadFeedback> {
2796        self.message_feedback.get(&message_id).copied()
2797    }
2798
2799    pub fn report_message_feedback(
2800        &mut self,
2801        message_id: MessageId,
2802        feedback: ThreadFeedback,
2803        cx: &mut Context<Self>,
2804    ) -> Task<Result<()>> {
2805        if self.message_feedback.get(&message_id) == Some(&feedback) {
2806            return Task::ready(Ok(()));
2807        }
2808
2809        let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2810        let serialized_thread = self.serialize(cx);
2811        let thread_id = self.id().clone();
2812        let client = self.project.read(cx).client();
2813
2814        let enabled_tool_names: Vec<String> = self
2815            .profile
2816            .enabled_tools(cx)
2817            .iter()
2818            .map(|(name, _)| name.clone().into())
2819            .collect();
2820
2821        self.message_feedback.insert(message_id, feedback);
2822
2823        cx.notify();
2824
2825        let message_content = self
2826            .message(message_id)
2827            .map(|msg| msg.to_string())
2828            .unwrap_or_default();
2829
2830        cx.background_spawn(async move {
2831            let final_project_snapshot = final_project_snapshot.await;
2832            let serialized_thread = serialized_thread.await?;
2833            let thread_data =
2834                serde_json::to_value(serialized_thread).unwrap_or_else(|_| serde_json::Value::Null);
2835
2836            let rating = match feedback {
2837                ThreadFeedback::Positive => "positive",
2838                ThreadFeedback::Negative => "negative",
2839            };
2840            telemetry::event!(
2841                "Assistant Thread Rated",
2842                rating,
2843                thread_id,
2844                enabled_tool_names,
2845                message_id = message_id.0,
2846                message_content,
2847                thread_data,
2848                final_project_snapshot
2849            );
2850            client.telemetry().flush_events().await;
2851
2852            Ok(())
2853        })
2854    }
2855
2856    pub fn report_feedback(
2857        &mut self,
2858        feedback: ThreadFeedback,
2859        cx: &mut Context<Self>,
2860    ) -> Task<Result<()>> {
2861        let last_assistant_message_id = self
2862            .messages
2863            .iter()
2864            .rev()
2865            .find(|msg| msg.role == Role::Assistant)
2866            .map(|msg| msg.id);
2867
2868        if let Some(message_id) = last_assistant_message_id {
2869            self.report_message_feedback(message_id, feedback, cx)
2870        } else {
2871            let final_project_snapshot = Self::project_snapshot(self.project.clone(), cx);
2872            let serialized_thread = self.serialize(cx);
2873            let thread_id = self.id().clone();
2874            let client = self.project.read(cx).client();
2875            self.feedback = Some(feedback);
2876            cx.notify();
2877
2878            cx.background_spawn(async move {
2879                let final_project_snapshot = final_project_snapshot.await;
2880                let serialized_thread = serialized_thread.await?;
2881                let thread_data = serde_json::to_value(serialized_thread)
2882                    .unwrap_or_else(|_| serde_json::Value::Null);
2883
2884                let rating = match feedback {
2885                    ThreadFeedback::Positive => "positive",
2886                    ThreadFeedback::Negative => "negative",
2887                };
2888                telemetry::event!(
2889                    "Assistant Thread Rated",
2890                    rating,
2891                    thread_id,
2892                    thread_data,
2893                    final_project_snapshot
2894                );
2895                client.telemetry().flush_events().await;
2896
2897                Ok(())
2898            })
2899        }
2900    }
2901
2902    /// Create a snapshot of the current project state including git information and unsaved buffers.
2903    fn project_snapshot(
2904        project: Entity<Project>,
2905        cx: &mut Context<Self>,
2906    ) -> Task<Arc<ProjectSnapshot>> {
2907        let git_store = project.read(cx).git_store().clone();
2908        let worktree_snapshots: Vec<_> = project
2909            .read(cx)
2910            .visible_worktrees(cx)
2911            .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
2912            .collect();
2913
2914        cx.spawn(async move |_, cx| {
2915            let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
2916
2917            let mut unsaved_buffers = Vec::new();
2918            cx.update(|app_cx| {
2919                let buffer_store = project.read(app_cx).buffer_store();
2920                for buffer_handle in buffer_store.read(app_cx).buffers() {
2921                    let buffer = buffer_handle.read(app_cx);
2922                    if buffer.is_dirty() {
2923                        if let Some(file) = buffer.file() {
2924                            let path = file.path().to_string_lossy().to_string();
2925                            unsaved_buffers.push(path);
2926                        }
2927                    }
2928                }
2929            })
2930            .ok();
2931
2932            Arc::new(ProjectSnapshot {
2933                worktree_snapshots,
2934                unsaved_buffer_paths: unsaved_buffers,
2935                timestamp: Utc::now(),
2936            })
2937        })
2938    }
2939
2940    fn worktree_snapshot(
2941        worktree: Entity<project::Worktree>,
2942        git_store: Entity<GitStore>,
2943        cx: &App,
2944    ) -> Task<WorktreeSnapshot> {
2945        cx.spawn(async move |cx| {
2946            // Get worktree path and snapshot
2947            let worktree_info = cx.update(|app_cx| {
2948                let worktree = worktree.read(app_cx);
2949                let path = worktree.abs_path().to_string_lossy().to_string();
2950                let snapshot = worktree.snapshot();
2951                (path, snapshot)
2952            });
2953
2954            let Ok((worktree_path, _snapshot)) = worktree_info else {
2955                return WorktreeSnapshot {
2956                    worktree_path: String::new(),
2957                    git_state: None,
2958                };
2959            };
2960
2961            let git_state = git_store
2962                .update(cx, |git_store, cx| {
2963                    git_store
2964                        .repositories()
2965                        .values()
2966                        .find(|repo| {
2967                            repo.read(cx)
2968                                .abs_path_to_repo_path(&worktree.read(cx).abs_path())
2969                                .is_some()
2970                        })
2971                        .cloned()
2972                })
2973                .ok()
2974                .flatten()
2975                .map(|repo| {
2976                    repo.update(cx, |repo, _| {
2977                        let current_branch =
2978                            repo.branch.as_ref().map(|branch| branch.name().to_owned());
2979                        repo.send_job(None, |state, _| async move {
2980                            let RepositoryState::Local { backend, .. } = state else {
2981                                return GitState {
2982                                    remote_url: None,
2983                                    head_sha: None,
2984                                    current_branch,
2985                                    diff: None,
2986                                };
2987                            };
2988
2989                            let remote_url = backend.remote_url("origin");
2990                            let head_sha = backend.head_sha().await;
2991                            let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
2992
2993                            GitState {
2994                                remote_url,
2995                                head_sha,
2996                                current_branch,
2997                                diff,
2998                            }
2999                        })
3000                    })
3001                });
3002
3003            let git_state = match git_state {
3004                Some(git_state) => match git_state.ok() {
3005                    Some(git_state) => git_state.await.ok(),
3006                    None => None,
3007                },
3008                None => None,
3009            };
3010
3011            WorktreeSnapshot {
3012                worktree_path,
3013                git_state,
3014            }
3015        })
3016    }
3017
3018    pub fn to_markdown(&self, cx: &App) -> Result<String> {
3019        let mut markdown = Vec::new();
3020
3021        let summary = self.summary().or_default();
3022        writeln!(markdown, "# {summary}\n")?;
3023
3024        for message in self.messages() {
3025            writeln!(
3026                markdown,
3027                "## {role}\n",
3028                role = match message.role {
3029                    Role::User => "User",
3030                    Role::Assistant => "Agent",
3031                    Role::System => "System",
3032                }
3033            )?;
3034
3035            if !message.loaded_context.text.is_empty() {
3036                writeln!(markdown, "{}", message.loaded_context.text)?;
3037            }
3038
3039            if !message.loaded_context.images.is_empty() {
3040                writeln!(
3041                    markdown,
3042                    "\n{} images attached as context.\n",
3043                    message.loaded_context.images.len()
3044                )?;
3045            }
3046
3047            for segment in &message.segments {
3048                match segment {
3049                    MessageSegment::Text(text) => writeln!(markdown, "{}\n", text)?,
3050                    MessageSegment::Thinking { text, .. } => {
3051                        writeln!(markdown, "<think>\n{}\n</think>\n", text)?
3052                    }
3053                    MessageSegment::RedactedThinking(_) => {}
3054                }
3055            }
3056
3057            for tool_use in self.tool_uses_for_message(message.id, cx) {
3058                writeln!(
3059                    markdown,
3060                    "**Use Tool: {} ({})**",
3061                    tool_use.name, tool_use.id
3062                )?;
3063                writeln!(markdown, "```json")?;
3064                writeln!(
3065                    markdown,
3066                    "{}",
3067                    serde_json::to_string_pretty(&tool_use.input)?
3068                )?;
3069                writeln!(markdown, "```")?;
3070            }
3071
3072            for tool_result in self.tool_results_for_message(message.id) {
3073                write!(markdown, "\n**Tool Results: {}", tool_result.tool_use_id)?;
3074                if tool_result.is_error {
3075                    write!(markdown, " (Error)")?;
3076                }
3077
3078                writeln!(markdown, "**\n")?;
3079                match &tool_result.content {
3080                    LanguageModelToolResultContent::Text(text) => {
3081                        writeln!(markdown, "{text}")?;
3082                    }
3083                    LanguageModelToolResultContent::Image(image) => {
3084                        writeln!(markdown, "![Image](data:base64,{})", image.source)?;
3085                    }
3086                }
3087
3088                if let Some(output) = tool_result.output.as_ref() {
3089                    writeln!(
3090                        markdown,
3091                        "\n\nDebug Output:\n\n```json\n{}\n```\n",
3092                        serde_json::to_string_pretty(output)?
3093                    )?;
3094                }
3095            }
3096        }
3097
3098        Ok(String::from_utf8_lossy(&markdown).to_string())
3099    }
3100
3101    pub fn keep_edits_in_range(
3102        &mut self,
3103        buffer: Entity<language::Buffer>,
3104        buffer_range: Range<language::Anchor>,
3105        cx: &mut Context<Self>,
3106    ) {
3107        self.action_log.update(cx, |action_log, cx| {
3108            action_log.keep_edits_in_range(buffer, buffer_range, cx)
3109        });
3110    }
3111
3112    pub fn keep_all_edits(&mut self, cx: &mut Context<Self>) {
3113        self.action_log
3114            .update(cx, |action_log, cx| action_log.keep_all_edits(cx));
3115    }
3116
3117    pub fn reject_edits_in_ranges(
3118        &mut self,
3119        buffer: Entity<language::Buffer>,
3120        buffer_ranges: Vec<Range<language::Anchor>>,
3121        cx: &mut Context<Self>,
3122    ) -> Task<Result<()>> {
3123        self.action_log.update(cx, |action_log, cx| {
3124            action_log.reject_edits_in_ranges(buffer, buffer_ranges, cx)
3125        })
3126    }
3127
3128    pub fn action_log(&self) -> &Entity<ActionLog> {
3129        &self.action_log
3130    }
3131
3132    pub fn project(&self) -> &Entity<Project> {
3133        &self.project
3134    }
3135
3136    pub fn auto_capture_telemetry(&mut self, cx: &mut Context<Self>) {
3137        if !cx.has_flag::<feature_flags::ThreadAutoCaptureFeatureFlag>() {
3138            return;
3139        }
3140
3141        let now = Instant::now();
3142        if let Some(last) = self.last_auto_capture_at {
3143            if now.duration_since(last).as_secs() < 10 {
3144                return;
3145            }
3146        }
3147
3148        self.last_auto_capture_at = Some(now);
3149
3150        let thread_id = self.id().clone();
3151        let github_login = self
3152            .project
3153            .read(cx)
3154            .user_store()
3155            .read(cx)
3156            .current_user()
3157            .map(|user| user.github_login.clone());
3158        let client = self.project.read(cx).client();
3159        let serialize_task = self.serialize(cx);
3160
3161        cx.background_executor()
3162            .spawn(async move {
3163                if let Ok(serialized_thread) = serialize_task.await {
3164                    if let Ok(thread_data) = serde_json::to_value(serialized_thread) {
3165                        telemetry::event!(
3166                            "Agent Thread Auto-Captured",
3167                            thread_id = thread_id.to_string(),
3168                            thread_data = thread_data,
3169                            auto_capture_reason = "tracked_user",
3170                            github_login = github_login
3171                        );
3172
3173                        client.telemetry().flush_events().await;
3174                    }
3175                }
3176            })
3177            .detach();
3178    }
3179
3180    pub fn cumulative_token_usage(&self) -> TokenUsage {
3181        self.cumulative_token_usage
3182    }
3183
3184    pub fn token_usage_up_to_message(&self, message_id: MessageId) -> TotalTokenUsage {
3185        let Some(model) = self.configured_model.as_ref() else {
3186            return TotalTokenUsage::default();
3187        };
3188
3189        let max = model
3190            .model
3191            .max_token_count_for_mode(self.completion_mode().into());
3192
3193        let index = self
3194            .messages
3195            .iter()
3196            .position(|msg| msg.id == message_id)
3197            .unwrap_or(0);
3198
3199        if index == 0 {
3200            return TotalTokenUsage { total: 0, max };
3201        }
3202
3203        let token_usage = &self
3204            .request_token_usage
3205            .get(index - 1)
3206            .cloned()
3207            .unwrap_or_default();
3208
3209        TotalTokenUsage {
3210            total: token_usage.total_tokens(),
3211            max,
3212        }
3213    }
3214
3215    pub fn total_token_usage(&self) -> Option<TotalTokenUsage> {
3216        let model = self.configured_model.as_ref()?;
3217
3218        let max = model
3219            .model
3220            .max_token_count_for_mode(self.completion_mode().into());
3221
3222        if let Some(exceeded_error) = &self.exceeded_window_error {
3223            if model.model.id() == exceeded_error.model_id {
3224                return Some(TotalTokenUsage {
3225                    total: exceeded_error.token_count,
3226                    max,
3227                });
3228            }
3229        }
3230
3231        let total = self
3232            .token_usage_at_last_message()
3233            .unwrap_or_default()
3234            .total_tokens();
3235
3236        Some(TotalTokenUsage { total, max })
3237    }
3238
3239    fn token_usage_at_last_message(&self) -> Option<TokenUsage> {
3240        self.request_token_usage
3241            .get(self.messages.len().saturating_sub(1))
3242            .or_else(|| self.request_token_usage.last())
3243            .cloned()
3244    }
3245
3246    fn update_token_usage_at_last_message(&mut self, token_usage: TokenUsage) {
3247        let placeholder = self.token_usage_at_last_message().unwrap_or_default();
3248        self.request_token_usage
3249            .resize(self.messages.len(), placeholder);
3250
3251        if let Some(last) = self.request_token_usage.last_mut() {
3252            *last = token_usage;
3253        }
3254    }
3255
3256    fn update_model_request_usage(&self, amount: u32, limit: UsageLimit, cx: &mut Context<Self>) {
3257        self.project
3258            .read(cx)
3259            .user_store()
3260            .update(cx, |user_store, cx| {
3261                user_store.update_model_request_usage(
3262                    ModelRequestUsage(RequestUsage {
3263                        amount: amount as i32,
3264                        limit,
3265                    }),
3266                    cx,
3267                )
3268            });
3269    }
3270
3271    pub fn deny_tool_use(
3272        &mut self,
3273        tool_use_id: LanguageModelToolUseId,
3274        tool_name: Arc<str>,
3275        window: Option<AnyWindowHandle>,
3276        cx: &mut Context<Self>,
3277    ) {
3278        let err = Err(anyhow::anyhow!(
3279            "Permission to run tool action denied by user"
3280        ));
3281
3282        self.tool_use.insert_tool_output(
3283            tool_use_id.clone(),
3284            tool_name,
3285            err,
3286            self.configured_model.as_ref(),
3287            self.completion_mode,
3288        );
3289        self.tool_finished(tool_use_id.clone(), None, true, window, cx);
3290    }
3291}
3292
3293#[derive(Debug, Clone, Error)]
3294pub enum ThreadError {
3295    #[error("Payment required")]
3296    PaymentRequired,
3297    #[error("Model request limit reached")]
3298    ModelRequestLimitReached { plan: Plan },
3299    #[error("Message {header}: {message}")]
3300    Message {
3301        header: SharedString,
3302        message: SharedString,
3303    },
3304    #[error("Retryable error: {message}")]
3305    RetryableError {
3306        message: SharedString,
3307        can_enable_burn_mode: bool,
3308    },
3309}
3310
3311#[derive(Debug, Clone)]
3312pub enum ThreadEvent {
3313    ShowError(ThreadError),
3314    StreamedCompletion,
3315    ReceivedTextChunk,
3316    NewRequest,
3317    StreamedAssistantText(MessageId, String),
3318    StreamedAssistantThinking(MessageId, String),
3319    StreamedToolUse {
3320        tool_use_id: LanguageModelToolUseId,
3321        ui_text: Arc<str>,
3322        input: serde_json::Value,
3323    },
3324    MissingToolUse {
3325        tool_use_id: LanguageModelToolUseId,
3326        ui_text: Arc<str>,
3327    },
3328    InvalidToolInput {
3329        tool_use_id: LanguageModelToolUseId,
3330        ui_text: Arc<str>,
3331        invalid_input_json: Arc<str>,
3332    },
3333    Stopped(Result<StopReason, Arc<anyhow::Error>>),
3334    MessageAdded(MessageId),
3335    MessageEdited(MessageId),
3336    MessageDeleted(MessageId),
3337    SummaryGenerated,
3338    SummaryChanged,
3339    UsePendingTools {
3340        tool_uses: Vec<PendingToolUse>,
3341    },
3342    ToolFinished {
3343        #[allow(unused)]
3344        tool_use_id: LanguageModelToolUseId,
3345        /// The pending tool use that corresponds to this tool.
3346        pending_tool_use: Option<PendingToolUse>,
3347    },
3348    CheckpointChanged,
3349    ToolConfirmationNeeded,
3350    ToolUseLimitReached,
3351    CancelEditing,
3352    CompletionCanceled,
3353    ProfileChanged,
3354}
3355
3356impl EventEmitter<ThreadEvent> for Thread {}
3357
3358struct PendingCompletion {
3359    id: usize,
3360    queue_state: QueueState,
3361    _task: Task<()>,
3362}
3363
3364#[cfg(test)]
3365mod tests {
3366    use super::*;
3367    use crate::{
3368        context::load_context, context_store::ContextStore, thread_store, thread_store::ThreadStore,
3369    };
3370
3371    // Test-specific constants
3372    const TEST_RATE_LIMIT_RETRY_SECS: u64 = 30;
3373    use agent_settings::{AgentProfileId, AgentSettings, LanguageModelParameters};
3374    use assistant_tool::ToolRegistry;
3375    use assistant_tools;
3376    use futures::StreamExt;
3377    use futures::future::BoxFuture;
3378    use futures::stream::BoxStream;
3379    use gpui::TestAppContext;
3380    use http_client;
3381    use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
3382    use language_model::{
3383        LanguageModelCompletionError, LanguageModelName, LanguageModelProviderId,
3384        LanguageModelProviderName, LanguageModelToolChoice,
3385    };
3386    use parking_lot::Mutex;
3387    use project::{FakeFs, Project};
3388    use prompt_store::PromptBuilder;
3389    use serde_json::json;
3390    use settings::{Settings, SettingsStore};
3391    use std::sync::Arc;
3392    use std::time::Duration;
3393    use theme::ThemeSettings;
3394    use util::path;
3395    use workspace::Workspace;
3396
3397    #[gpui::test]
3398    async fn test_message_with_context(cx: &mut TestAppContext) {
3399        init_test_settings(cx);
3400
3401        let project = create_test_project(
3402            cx,
3403            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3404        )
3405        .await;
3406
3407        let (_workspace, _thread_store, thread, context_store, model) =
3408            setup_test_environment(cx, project.clone()).await;
3409
3410        add_file_to_context(&project, &context_store, "test/code.rs", cx)
3411            .await
3412            .unwrap();
3413
3414        let context =
3415            context_store.read_with(cx, |store, _| store.context().next().cloned().unwrap());
3416        let loaded_context = cx
3417            .update(|cx| load_context(vec![context], &project, &None, cx))
3418            .await;
3419
3420        // Insert user message with context
3421        let message_id = thread.update(cx, |thread, cx| {
3422            thread.insert_user_message(
3423                "Please explain this code",
3424                loaded_context,
3425                None,
3426                Vec::new(),
3427                cx,
3428            )
3429        });
3430
3431        // Check content and context in message object
3432        let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3433
3434        // Use different path format strings based on platform for the test
3435        #[cfg(windows)]
3436        let path_part = r"test\code.rs";
3437        #[cfg(not(windows))]
3438        let path_part = "test/code.rs";
3439
3440        let expected_context = format!(
3441            r#"
3442<context>
3443The following items were attached by the user. They are up-to-date and don't need to be re-read.
3444
3445<files>
3446```rs {path_part}
3447fn main() {{
3448    println!("Hello, world!");
3449}}
3450```
3451</files>
3452</context>
3453"#
3454        );
3455
3456        assert_eq!(message.role, Role::User);
3457        assert_eq!(message.segments.len(), 1);
3458        assert_eq!(
3459            message.segments[0],
3460            MessageSegment::Text("Please explain this code".to_string())
3461        );
3462        assert_eq!(message.loaded_context.text, expected_context);
3463
3464        // Check message in request
3465        let request = thread.update(cx, |thread, cx| {
3466            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3467        });
3468
3469        assert_eq!(request.messages.len(), 2);
3470        let expected_full_message = format!("{}Please explain this code", expected_context);
3471        assert_eq!(request.messages[1].string_contents(), expected_full_message);
3472    }
3473
3474    #[gpui::test]
3475    async fn test_only_include_new_contexts(cx: &mut TestAppContext) {
3476        init_test_settings(cx);
3477
3478        let project = create_test_project(
3479            cx,
3480            json!({
3481                "file1.rs": "fn function1() {}\n",
3482                "file2.rs": "fn function2() {}\n",
3483                "file3.rs": "fn function3() {}\n",
3484                "file4.rs": "fn function4() {}\n",
3485            }),
3486        )
3487        .await;
3488
3489        let (_, _thread_store, thread, context_store, model) =
3490            setup_test_environment(cx, project.clone()).await;
3491
3492        // First message with context 1
3493        add_file_to_context(&project, &context_store, "test/file1.rs", cx)
3494            .await
3495            .unwrap();
3496        let new_contexts = context_store.update(cx, |store, cx| {
3497            store.new_context_for_thread(thread.read(cx), None)
3498        });
3499        assert_eq!(new_contexts.len(), 1);
3500        let loaded_context = cx
3501            .update(|cx| load_context(new_contexts, &project, &None, cx))
3502            .await;
3503        let message1_id = thread.update(cx, |thread, cx| {
3504            thread.insert_user_message("Message 1", loaded_context, None, Vec::new(), cx)
3505        });
3506
3507        // Second message with contexts 1 and 2 (context 1 should be skipped as it's already included)
3508        add_file_to_context(&project, &context_store, "test/file2.rs", cx)
3509            .await
3510            .unwrap();
3511        let new_contexts = context_store.update(cx, |store, cx| {
3512            store.new_context_for_thread(thread.read(cx), None)
3513        });
3514        assert_eq!(new_contexts.len(), 1);
3515        let loaded_context = cx
3516            .update(|cx| load_context(new_contexts, &project, &None, cx))
3517            .await;
3518        let message2_id = thread.update(cx, |thread, cx| {
3519            thread.insert_user_message("Message 2", loaded_context, None, Vec::new(), cx)
3520        });
3521
3522        // Third message with all three contexts (contexts 1 and 2 should be skipped)
3523        //
3524        add_file_to_context(&project, &context_store, "test/file3.rs", cx)
3525            .await
3526            .unwrap();
3527        let new_contexts = context_store.update(cx, |store, cx| {
3528            store.new_context_for_thread(thread.read(cx), None)
3529        });
3530        assert_eq!(new_contexts.len(), 1);
3531        let loaded_context = cx
3532            .update(|cx| load_context(new_contexts, &project, &None, cx))
3533            .await;
3534        let message3_id = thread.update(cx, |thread, cx| {
3535            thread.insert_user_message("Message 3", loaded_context, None, Vec::new(), cx)
3536        });
3537
3538        // Check what contexts are included in each message
3539        let (message1, message2, message3) = thread.read_with(cx, |thread, _| {
3540            (
3541                thread.message(message1_id).unwrap().clone(),
3542                thread.message(message2_id).unwrap().clone(),
3543                thread.message(message3_id).unwrap().clone(),
3544            )
3545        });
3546
3547        // First message should include context 1
3548        assert!(message1.loaded_context.text.contains("file1.rs"));
3549
3550        // Second message should include only context 2 (not 1)
3551        assert!(!message2.loaded_context.text.contains("file1.rs"));
3552        assert!(message2.loaded_context.text.contains("file2.rs"));
3553
3554        // Third message should include only context 3 (not 1 or 2)
3555        assert!(!message3.loaded_context.text.contains("file1.rs"));
3556        assert!(!message3.loaded_context.text.contains("file2.rs"));
3557        assert!(message3.loaded_context.text.contains("file3.rs"));
3558
3559        // Check entire request to make sure all contexts are properly included
3560        let request = thread.update(cx, |thread, cx| {
3561            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3562        });
3563
3564        // The request should contain all 3 messages
3565        assert_eq!(request.messages.len(), 4);
3566
3567        // Check that the contexts are properly formatted in each message
3568        assert!(request.messages[1].string_contents().contains("file1.rs"));
3569        assert!(!request.messages[1].string_contents().contains("file2.rs"));
3570        assert!(!request.messages[1].string_contents().contains("file3.rs"));
3571
3572        assert!(!request.messages[2].string_contents().contains("file1.rs"));
3573        assert!(request.messages[2].string_contents().contains("file2.rs"));
3574        assert!(!request.messages[2].string_contents().contains("file3.rs"));
3575
3576        assert!(!request.messages[3].string_contents().contains("file1.rs"));
3577        assert!(!request.messages[3].string_contents().contains("file2.rs"));
3578        assert!(request.messages[3].string_contents().contains("file3.rs"));
3579
3580        add_file_to_context(&project, &context_store, "test/file4.rs", cx)
3581            .await
3582            .unwrap();
3583        let new_contexts = context_store.update(cx, |store, cx| {
3584            store.new_context_for_thread(thread.read(cx), Some(message2_id))
3585        });
3586        assert_eq!(new_contexts.len(), 3);
3587        let loaded_context = cx
3588            .update(|cx| load_context(new_contexts, &project, &None, cx))
3589            .await
3590            .loaded_context;
3591
3592        assert!(!loaded_context.text.contains("file1.rs"));
3593        assert!(loaded_context.text.contains("file2.rs"));
3594        assert!(loaded_context.text.contains("file3.rs"));
3595        assert!(loaded_context.text.contains("file4.rs"));
3596
3597        let new_contexts = context_store.update(cx, |store, cx| {
3598            // Remove file4.rs
3599            store.remove_context(&loaded_context.contexts[2].handle(), cx);
3600            store.new_context_for_thread(thread.read(cx), Some(message2_id))
3601        });
3602        assert_eq!(new_contexts.len(), 2);
3603        let loaded_context = cx
3604            .update(|cx| load_context(new_contexts, &project, &None, cx))
3605            .await
3606            .loaded_context;
3607
3608        assert!(!loaded_context.text.contains("file1.rs"));
3609        assert!(loaded_context.text.contains("file2.rs"));
3610        assert!(loaded_context.text.contains("file3.rs"));
3611        assert!(!loaded_context.text.contains("file4.rs"));
3612
3613        let new_contexts = context_store.update(cx, |store, cx| {
3614            // Remove file3.rs
3615            store.remove_context(&loaded_context.contexts[1].handle(), cx);
3616            store.new_context_for_thread(thread.read(cx), Some(message2_id))
3617        });
3618        assert_eq!(new_contexts.len(), 1);
3619        let loaded_context = cx
3620            .update(|cx| load_context(new_contexts, &project, &None, cx))
3621            .await
3622            .loaded_context;
3623
3624        assert!(!loaded_context.text.contains("file1.rs"));
3625        assert!(loaded_context.text.contains("file2.rs"));
3626        assert!(!loaded_context.text.contains("file3.rs"));
3627        assert!(!loaded_context.text.contains("file4.rs"));
3628    }
3629
3630    #[gpui::test]
3631    async fn test_message_without_files(cx: &mut TestAppContext) {
3632        init_test_settings(cx);
3633
3634        let project = create_test_project(
3635            cx,
3636            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3637        )
3638        .await;
3639
3640        let (_, _thread_store, thread, _context_store, model) =
3641            setup_test_environment(cx, project.clone()).await;
3642
3643        // Insert user message without any context (empty context vector)
3644        let message_id = thread.update(cx, |thread, cx| {
3645            thread.insert_user_message(
3646                "What is the best way to learn Rust?",
3647                ContextLoadResult::default(),
3648                None,
3649                Vec::new(),
3650                cx,
3651            )
3652        });
3653
3654        // Check content and context in message object
3655        let message = thread.read_with(cx, |thread, _| thread.message(message_id).unwrap().clone());
3656
3657        // Context should be empty when no files are included
3658        assert_eq!(message.role, Role::User);
3659        assert_eq!(message.segments.len(), 1);
3660        assert_eq!(
3661            message.segments[0],
3662            MessageSegment::Text("What is the best way to learn Rust?".to_string())
3663        );
3664        assert_eq!(message.loaded_context.text, "");
3665
3666        // Check message in request
3667        let request = thread.update(cx, |thread, cx| {
3668            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3669        });
3670
3671        assert_eq!(request.messages.len(), 2);
3672        assert_eq!(
3673            request.messages[1].string_contents(),
3674            "What is the best way to learn Rust?"
3675        );
3676
3677        // Add second message, also without context
3678        let message2_id = thread.update(cx, |thread, cx| {
3679            thread.insert_user_message(
3680                "Are there any good books?",
3681                ContextLoadResult::default(),
3682                None,
3683                Vec::new(),
3684                cx,
3685            )
3686        });
3687
3688        let message2 =
3689            thread.read_with(cx, |thread, _| thread.message(message2_id).unwrap().clone());
3690        assert_eq!(message2.loaded_context.text, "");
3691
3692        // Check that both messages appear in the request
3693        let request = thread.update(cx, |thread, cx| {
3694            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3695        });
3696
3697        assert_eq!(request.messages.len(), 3);
3698        assert_eq!(
3699            request.messages[1].string_contents(),
3700            "What is the best way to learn Rust?"
3701        );
3702        assert_eq!(
3703            request.messages[2].string_contents(),
3704            "Are there any good books?"
3705        );
3706    }
3707
3708    #[gpui::test]
3709    #[ignore] // turn this test on when project_notifications tool is re-enabled
3710    async fn test_stale_buffer_notification(cx: &mut TestAppContext) {
3711        init_test_settings(cx);
3712
3713        let project = create_test_project(
3714            cx,
3715            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3716        )
3717        .await;
3718
3719        let (_workspace, _thread_store, thread, context_store, model) =
3720            setup_test_environment(cx, project.clone()).await;
3721
3722        // Add a buffer to the context. This will be a tracked buffer
3723        let buffer = add_file_to_context(&project, &context_store, "test/code.rs", cx)
3724            .await
3725            .unwrap();
3726
3727        let context = context_store
3728            .read_with(cx, |store, _| store.context().next().cloned())
3729            .unwrap();
3730        let loaded_context = cx
3731            .update(|cx| load_context(vec![context], &project, &None, cx))
3732            .await;
3733
3734        // Insert user message and assistant response
3735        thread.update(cx, |thread, cx| {
3736            thread.insert_user_message("Explain this code", loaded_context, None, Vec::new(), cx);
3737            thread.insert_assistant_message(
3738                vec![MessageSegment::Text("This code prints 42.".into())],
3739                cx,
3740            );
3741        });
3742        cx.run_until_parked();
3743
3744        // We shouldn't have a stale buffer notification yet
3745        let notifications = thread.read_with(cx, |thread, _| {
3746            find_tool_uses(thread, "project_notifications")
3747        });
3748        assert!(
3749            notifications.is_empty(),
3750            "Should not have stale buffer notification before buffer is modified"
3751        );
3752
3753        // Modify the buffer
3754        buffer.update(cx, |buffer, cx| {
3755            buffer.edit(
3756                [(1..1, "\n    println!(\"Added a new line\");\n")],
3757                None,
3758                cx,
3759            );
3760        });
3761
3762        // Insert another user message
3763        thread.update(cx, |thread, cx| {
3764            thread.insert_user_message(
3765                "What does the code do now?",
3766                ContextLoadResult::default(),
3767                None,
3768                Vec::new(),
3769                cx,
3770            )
3771        });
3772        cx.run_until_parked();
3773
3774        // Check for the stale buffer warning
3775        thread.update(cx, |thread, cx| {
3776            thread.flush_notifications(model.clone(), CompletionIntent::UserPrompt, cx)
3777        });
3778        cx.run_until_parked();
3779
3780        let notifications = thread.read_with(cx, |thread, _cx| {
3781            find_tool_uses(thread, "project_notifications")
3782        });
3783
3784        let [notification] = notifications.as_slice() else {
3785            panic!("Should have a `project_notifications` tool use");
3786        };
3787
3788        let Some(notification_content) = notification.content.to_str() else {
3789            panic!("`project_notifications` should return text");
3790        };
3791
3792        assert!(notification_content.contains("These files have changed since the last read:"));
3793        assert!(notification_content.contains("code.rs"));
3794
3795        // Insert another user message and flush notifications again
3796        thread.update(cx, |thread, cx| {
3797            thread.insert_user_message(
3798                "Can you tell me more?",
3799                ContextLoadResult::default(),
3800                None,
3801                Vec::new(),
3802                cx,
3803            )
3804        });
3805
3806        thread.update(cx, |thread, cx| {
3807            thread.flush_notifications(model.clone(), CompletionIntent::UserPrompt, cx)
3808        });
3809        cx.run_until_parked();
3810
3811        // There should be no new notifications (we already flushed one)
3812        let notifications = thread.read_with(cx, |thread, _cx| {
3813            find_tool_uses(thread, "project_notifications")
3814        });
3815
3816        assert_eq!(
3817            notifications.len(),
3818            1,
3819            "Should still have only one notification after second flush - no duplicates"
3820        );
3821    }
3822
3823    fn find_tool_uses(thread: &Thread, tool_name: &str) -> Vec<LanguageModelToolResult> {
3824        thread
3825            .messages()
3826            .flat_map(|message| {
3827                thread
3828                    .tool_results_for_message(message.id)
3829                    .into_iter()
3830                    .filter(|result| result.tool_name == tool_name.into())
3831                    .cloned()
3832                    .collect::<Vec<_>>()
3833            })
3834            .collect()
3835    }
3836
3837    #[gpui::test]
3838    async fn test_storing_profile_setting_per_thread(cx: &mut TestAppContext) {
3839        init_test_settings(cx);
3840
3841        let project = create_test_project(
3842            cx,
3843            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3844        )
3845        .await;
3846
3847        let (_workspace, thread_store, thread, _context_store, _model) =
3848            setup_test_environment(cx, project.clone()).await;
3849
3850        // Check that we are starting with the default profile
3851        let profile = cx.read(|cx| thread.read(cx).profile.clone());
3852        let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3853        assert_eq!(
3854            profile,
3855            AgentProfile::new(AgentProfileId::default(), tool_set)
3856        );
3857    }
3858
3859    #[gpui::test]
3860    async fn test_serializing_thread_profile(cx: &mut TestAppContext) {
3861        init_test_settings(cx);
3862
3863        let project = create_test_project(
3864            cx,
3865            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3866        )
3867        .await;
3868
3869        let (_workspace, thread_store, thread, _context_store, _model) =
3870            setup_test_environment(cx, project.clone()).await;
3871
3872        // Profile gets serialized with default values
3873        let serialized = thread
3874            .update(cx, |thread, cx| thread.serialize(cx))
3875            .await
3876            .unwrap();
3877
3878        assert_eq!(serialized.profile, Some(AgentProfileId::default()));
3879
3880        let deserialized = cx.update(|cx| {
3881            thread.update(cx, |thread, cx| {
3882                Thread::deserialize(
3883                    thread.id.clone(),
3884                    serialized,
3885                    thread.project.clone(),
3886                    thread.tools.clone(),
3887                    thread.prompt_builder.clone(),
3888                    thread.project_context.clone(),
3889                    None,
3890                    cx,
3891                )
3892            })
3893        });
3894        let tool_set = cx.read(|cx| thread_store.read(cx).tools());
3895
3896        assert_eq!(
3897            deserialized.profile,
3898            AgentProfile::new(AgentProfileId::default(), tool_set)
3899        );
3900    }
3901
3902    #[gpui::test]
3903    async fn test_temperature_setting(cx: &mut TestAppContext) {
3904        init_test_settings(cx);
3905
3906        let project = create_test_project(
3907            cx,
3908            json!({"code.rs": "fn main() {\n    println!(\"Hello, world!\");\n}"}),
3909        )
3910        .await;
3911
3912        let (_workspace, _thread_store, thread, _context_store, model) =
3913            setup_test_environment(cx, project.clone()).await;
3914
3915        // Both model and provider
3916        cx.update(|cx| {
3917            AgentSettings::override_global(
3918                AgentSettings {
3919                    model_parameters: vec![LanguageModelParameters {
3920                        provider: Some(model.provider_id().0.to_string().into()),
3921                        model: Some(model.id().0.clone()),
3922                        temperature: Some(0.66),
3923                    }],
3924                    ..AgentSettings::get_global(cx).clone()
3925                },
3926                cx,
3927            );
3928        });
3929
3930        let request = thread.update(cx, |thread, cx| {
3931            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3932        });
3933        assert_eq!(request.temperature, Some(0.66));
3934
3935        // Only model
3936        cx.update(|cx| {
3937            AgentSettings::override_global(
3938                AgentSettings {
3939                    model_parameters: vec![LanguageModelParameters {
3940                        provider: None,
3941                        model: Some(model.id().0.clone()),
3942                        temperature: Some(0.66),
3943                    }],
3944                    ..AgentSettings::get_global(cx).clone()
3945                },
3946                cx,
3947            );
3948        });
3949
3950        let request = thread.update(cx, |thread, cx| {
3951            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3952        });
3953        assert_eq!(request.temperature, Some(0.66));
3954
3955        // Only provider
3956        cx.update(|cx| {
3957            AgentSettings::override_global(
3958                AgentSettings {
3959                    model_parameters: vec![LanguageModelParameters {
3960                        provider: Some(model.provider_id().0.to_string().into()),
3961                        model: None,
3962                        temperature: Some(0.66),
3963                    }],
3964                    ..AgentSettings::get_global(cx).clone()
3965                },
3966                cx,
3967            );
3968        });
3969
3970        let request = thread.update(cx, |thread, cx| {
3971            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3972        });
3973        assert_eq!(request.temperature, Some(0.66));
3974
3975        // Same model name, different provider
3976        cx.update(|cx| {
3977            AgentSettings::override_global(
3978                AgentSettings {
3979                    model_parameters: vec![LanguageModelParameters {
3980                        provider: Some("anthropic".into()),
3981                        model: Some(model.id().0.clone()),
3982                        temperature: Some(0.66),
3983                    }],
3984                    ..AgentSettings::get_global(cx).clone()
3985                },
3986                cx,
3987            );
3988        });
3989
3990        let request = thread.update(cx, |thread, cx| {
3991            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
3992        });
3993        assert_eq!(request.temperature, None);
3994    }
3995
3996    #[gpui::test]
3997    async fn test_thread_summary(cx: &mut TestAppContext) {
3998        init_test_settings(cx);
3999
4000        let project = create_test_project(cx, json!({})).await;
4001
4002        let (_, _thread_store, thread, _context_store, model) =
4003            setup_test_environment(cx, project.clone()).await;
4004
4005        // Initial state should be pending
4006        thread.read_with(cx, |thread, _| {
4007            assert!(matches!(thread.summary(), ThreadSummary::Pending));
4008            assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4009        });
4010
4011        // Manually setting the summary should not be allowed in this state
4012        thread.update(cx, |thread, cx| {
4013            thread.set_summary("This should not work", cx);
4014        });
4015
4016        thread.read_with(cx, |thread, _| {
4017            assert!(matches!(thread.summary(), ThreadSummary::Pending));
4018        });
4019
4020        // Send a message
4021        thread.update(cx, |thread, cx| {
4022            thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
4023            thread.send_to_model(
4024                model.clone(),
4025                CompletionIntent::ThreadSummarization,
4026                None,
4027                cx,
4028            );
4029        });
4030
4031        let fake_model = model.as_fake();
4032        simulate_successful_response(&fake_model, cx);
4033
4034        // Should start generating summary when there are >= 2 messages
4035        thread.read_with(cx, |thread, _| {
4036            assert_eq!(*thread.summary(), ThreadSummary::Generating);
4037        });
4038
4039        // Should not be able to set the summary while generating
4040        thread.update(cx, |thread, cx| {
4041            thread.set_summary("This should not work either", cx);
4042        });
4043
4044        thread.read_with(cx, |thread, _| {
4045            assert!(matches!(thread.summary(), ThreadSummary::Generating));
4046            assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4047        });
4048
4049        cx.run_until_parked();
4050        fake_model.stream_last_completion_response("Brief");
4051        fake_model.stream_last_completion_response(" Introduction");
4052        fake_model.end_last_completion_stream();
4053        cx.run_until_parked();
4054
4055        // Summary should be set
4056        thread.read_with(cx, |thread, _| {
4057            assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4058            assert_eq!(thread.summary().or_default(), "Brief Introduction");
4059        });
4060
4061        // Now we should be able to set a summary
4062        thread.update(cx, |thread, cx| {
4063            thread.set_summary("Brief Intro", cx);
4064        });
4065
4066        thread.read_with(cx, |thread, _| {
4067            assert_eq!(thread.summary().or_default(), "Brief Intro");
4068        });
4069
4070        // Test setting an empty summary (should default to DEFAULT)
4071        thread.update(cx, |thread, cx| {
4072            thread.set_summary("", cx);
4073        });
4074
4075        thread.read_with(cx, |thread, _| {
4076            assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4077            assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
4078        });
4079    }
4080
4081    #[gpui::test]
4082    async fn test_thread_summary_error_set_manually(cx: &mut TestAppContext) {
4083        init_test_settings(cx);
4084
4085        let project = create_test_project(cx, json!({})).await;
4086
4087        let (_, _thread_store, thread, _context_store, model) =
4088            setup_test_environment(cx, project.clone()).await;
4089
4090        test_summarize_error(&model, &thread, cx);
4091
4092        // Now we should be able to set a summary
4093        thread.update(cx, |thread, cx| {
4094            thread.set_summary("Brief Intro", cx);
4095        });
4096
4097        thread.read_with(cx, |thread, _| {
4098            assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4099            assert_eq!(thread.summary().or_default(), "Brief Intro");
4100        });
4101    }
4102
4103    #[gpui::test]
4104    async fn test_thread_summary_error_retry(cx: &mut TestAppContext) {
4105        init_test_settings(cx);
4106
4107        let project = create_test_project(cx, json!({})).await;
4108
4109        let (_, _thread_store, thread, _context_store, model) =
4110            setup_test_environment(cx, project.clone()).await;
4111
4112        test_summarize_error(&model, &thread, cx);
4113
4114        // Sending another message should not trigger another summarize request
4115        thread.update(cx, |thread, cx| {
4116            thread.insert_user_message(
4117                "How are you?",
4118                ContextLoadResult::default(),
4119                None,
4120                vec![],
4121                cx,
4122            );
4123            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4124        });
4125
4126        let fake_model = model.as_fake();
4127        simulate_successful_response(&fake_model, cx);
4128
4129        thread.read_with(cx, |thread, _| {
4130            // State is still Error, not Generating
4131            assert!(matches!(thread.summary(), ThreadSummary::Error));
4132        });
4133
4134        // But the summarize request can be invoked manually
4135        thread.update(cx, |thread, cx| {
4136            thread.summarize(cx);
4137        });
4138
4139        thread.read_with(cx, |thread, _| {
4140            assert!(matches!(thread.summary(), ThreadSummary::Generating));
4141        });
4142
4143        cx.run_until_parked();
4144        fake_model.stream_last_completion_response("A successful summary");
4145        fake_model.end_last_completion_stream();
4146        cx.run_until_parked();
4147
4148        thread.read_with(cx, |thread, _| {
4149            assert!(matches!(thread.summary(), ThreadSummary::Ready(_)));
4150            assert_eq!(thread.summary().or_default(), "A successful summary");
4151        });
4152    }
4153
4154    // Helper to create a model that returns errors
4155    enum TestError {
4156        Overloaded,
4157        InternalServerError,
4158    }
4159
4160    struct ErrorInjector {
4161        inner: Arc<FakeLanguageModel>,
4162        error_type: TestError,
4163    }
4164
4165    impl ErrorInjector {
4166        fn new(error_type: TestError) -> Self {
4167            Self {
4168                inner: Arc::new(FakeLanguageModel::default()),
4169                error_type,
4170            }
4171        }
4172    }
4173
4174    impl LanguageModel for ErrorInjector {
4175        fn id(&self) -> LanguageModelId {
4176            self.inner.id()
4177        }
4178
4179        fn name(&self) -> LanguageModelName {
4180            self.inner.name()
4181        }
4182
4183        fn provider_id(&self) -> LanguageModelProviderId {
4184            self.inner.provider_id()
4185        }
4186
4187        fn provider_name(&self) -> LanguageModelProviderName {
4188            self.inner.provider_name()
4189        }
4190
4191        fn supports_tools(&self) -> bool {
4192            self.inner.supports_tools()
4193        }
4194
4195        fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4196            self.inner.supports_tool_choice(choice)
4197        }
4198
4199        fn supports_images(&self) -> bool {
4200            self.inner.supports_images()
4201        }
4202
4203        fn telemetry_id(&self) -> String {
4204            self.inner.telemetry_id()
4205        }
4206
4207        fn max_token_count(&self) -> u64 {
4208            self.inner.max_token_count()
4209        }
4210
4211        fn count_tokens(
4212            &self,
4213            request: LanguageModelRequest,
4214            cx: &App,
4215        ) -> BoxFuture<'static, Result<u64>> {
4216            self.inner.count_tokens(request, cx)
4217        }
4218
4219        fn stream_completion(
4220            &self,
4221            _request: LanguageModelRequest,
4222            _cx: &AsyncApp,
4223        ) -> BoxFuture<
4224            'static,
4225            Result<
4226                BoxStream<
4227                    'static,
4228                    Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4229                >,
4230                LanguageModelCompletionError,
4231            >,
4232        > {
4233            let error = match self.error_type {
4234                TestError::Overloaded => LanguageModelCompletionError::ServerOverloaded {
4235                    provider: self.provider_name(),
4236                    retry_after: None,
4237                },
4238                TestError::InternalServerError => {
4239                    LanguageModelCompletionError::ApiInternalServerError {
4240                        provider: self.provider_name(),
4241                        message: "I'm a teapot orbiting the sun".to_string(),
4242                    }
4243                }
4244            };
4245            async move {
4246                let stream = futures::stream::once(async move { Err(error) });
4247                Ok(stream.boxed())
4248            }
4249            .boxed()
4250        }
4251
4252        fn as_fake(&self) -> &FakeLanguageModel {
4253            &self.inner
4254        }
4255    }
4256
4257    #[gpui::test]
4258    async fn test_retry_on_overloaded_error(cx: &mut TestAppContext) {
4259        init_test_settings(cx);
4260
4261        let project = create_test_project(cx, json!({})).await;
4262        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4263
4264        // Enable Burn Mode to allow retries
4265        thread.update(cx, |thread, _| {
4266            thread.set_completion_mode(CompletionMode::Burn);
4267        });
4268
4269        // Create model that returns overloaded error
4270        let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
4271
4272        // Insert a user message
4273        thread.update(cx, |thread, cx| {
4274            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4275        });
4276
4277        // Start completion
4278        thread.update(cx, |thread, cx| {
4279            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4280        });
4281
4282        cx.run_until_parked();
4283
4284        thread.read_with(cx, |thread, _| {
4285            assert!(thread.retry_state.is_some(), "Should have retry state");
4286            let retry_state = thread.retry_state.as_ref().unwrap();
4287            assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4288            assert_eq!(
4289                retry_state.max_attempts, MAX_RETRY_ATTEMPTS,
4290                "Should retry MAX_RETRY_ATTEMPTS times for overloaded errors"
4291            );
4292        });
4293
4294        // Check that a retry message was added
4295        thread.read_with(cx, |thread, _| {
4296            let mut messages = thread.messages();
4297            assert!(
4298                messages.any(|msg| {
4299                    msg.role == Role::System
4300                        && msg.ui_only
4301                        && msg.segments.iter().any(|seg| {
4302                            if let MessageSegment::Text(text) = seg {
4303                                text.contains("overloaded")
4304                                    && text
4305                                        .contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS))
4306                            } else {
4307                                false
4308                            }
4309                        })
4310                }),
4311                "Should have added a system retry message"
4312            );
4313        });
4314
4315        let retry_count = thread.update(cx, |thread, _| {
4316            thread
4317                .messages
4318                .iter()
4319                .filter(|m| {
4320                    m.ui_only
4321                        && m.segments.iter().any(|s| {
4322                            if let MessageSegment::Text(text) = s {
4323                                text.contains("Retrying") && text.contains("seconds")
4324                            } else {
4325                                false
4326                            }
4327                        })
4328                })
4329                .count()
4330        });
4331
4332        assert_eq!(retry_count, 1, "Should have one retry message");
4333    }
4334
4335    #[gpui::test]
4336    async fn test_retry_on_internal_server_error(cx: &mut TestAppContext) {
4337        init_test_settings(cx);
4338
4339        let project = create_test_project(cx, json!({})).await;
4340        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4341
4342        // Enable Burn Mode to allow retries
4343        thread.update(cx, |thread, _| {
4344            thread.set_completion_mode(CompletionMode::Burn);
4345        });
4346
4347        // Create model that returns internal server error
4348        let model = Arc::new(ErrorInjector::new(TestError::InternalServerError));
4349
4350        // Insert a user message
4351        thread.update(cx, |thread, cx| {
4352            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4353        });
4354
4355        // Start completion
4356        thread.update(cx, |thread, cx| {
4357            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4358        });
4359
4360        cx.run_until_parked();
4361
4362        // Check retry state on thread
4363        thread.read_with(cx, |thread, _| {
4364            assert!(thread.retry_state.is_some(), "Should have retry state");
4365            let retry_state = thread.retry_state.as_ref().unwrap();
4366            assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4367            assert_eq!(
4368                retry_state.max_attempts, 3,
4369                "Should have correct max attempts"
4370            );
4371        });
4372
4373        // Check that a retry message was added with provider name
4374        thread.read_with(cx, |thread, _| {
4375            let mut messages = thread.messages();
4376            assert!(
4377                messages.any(|msg| {
4378                    msg.role == Role::System
4379                        && msg.ui_only
4380                        && msg.segments.iter().any(|seg| {
4381                            if let MessageSegment::Text(text) = seg {
4382                                text.contains("internal")
4383                                    && text.contains("Fake")
4384                                    && text.contains("Retrying")
4385                                    && text.contains("attempt 1 of 3")
4386                                    && text.contains("seconds")
4387                            } else {
4388                                false
4389                            }
4390                        })
4391                }),
4392                "Should have added a system retry message with provider name"
4393            );
4394        });
4395
4396        // Count retry messages
4397        let retry_count = thread.update(cx, |thread, _| {
4398            thread
4399                .messages
4400                .iter()
4401                .filter(|m| {
4402                    m.ui_only
4403                        && m.segments.iter().any(|s| {
4404                            if let MessageSegment::Text(text) = s {
4405                                text.contains("Retrying") && text.contains("seconds")
4406                            } else {
4407                                false
4408                            }
4409                        })
4410                })
4411                .count()
4412        });
4413
4414        assert_eq!(retry_count, 1, "Should have one retry message");
4415    }
4416
4417    #[gpui::test]
4418    async fn test_exponential_backoff_on_retries(cx: &mut TestAppContext) {
4419        init_test_settings(cx);
4420
4421        let project = create_test_project(cx, json!({})).await;
4422        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4423
4424        // Enable Burn Mode to allow retries
4425        thread.update(cx, |thread, _| {
4426            thread.set_completion_mode(CompletionMode::Burn);
4427        });
4428
4429        // Create model that returns internal server error
4430        let model = Arc::new(ErrorInjector::new(TestError::InternalServerError));
4431
4432        // Insert a user message
4433        thread.update(cx, |thread, cx| {
4434            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4435        });
4436
4437        // Track retry events and completion count
4438        // Track completion events
4439        let completion_count = Arc::new(Mutex::new(0));
4440        let completion_count_clone = completion_count.clone();
4441
4442        let _subscription = thread.update(cx, |_, cx| {
4443            cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4444                if let ThreadEvent::NewRequest = event {
4445                    *completion_count_clone.lock() += 1;
4446                }
4447            })
4448        });
4449
4450        // First attempt
4451        thread.update(cx, |thread, cx| {
4452            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4453        });
4454        cx.run_until_parked();
4455
4456        // Should have scheduled first retry - count retry messages
4457        let retry_count = thread.update(cx, |thread, _| {
4458            thread
4459                .messages
4460                .iter()
4461                .filter(|m| {
4462                    m.ui_only
4463                        && m.segments.iter().any(|s| {
4464                            if let MessageSegment::Text(text) = s {
4465                                text.contains("Retrying") && text.contains("seconds")
4466                            } else {
4467                                false
4468                            }
4469                        })
4470                })
4471                .count()
4472        });
4473        assert_eq!(retry_count, 1, "Should have scheduled first retry");
4474
4475        // Check retry state
4476        thread.read_with(cx, |thread, _| {
4477            assert!(thread.retry_state.is_some(), "Should have retry state");
4478            let retry_state = thread.retry_state.as_ref().unwrap();
4479            assert_eq!(retry_state.attempt, 1, "Should be first retry attempt");
4480            assert_eq!(
4481                retry_state.max_attempts, 3,
4482                "Internal server errors should retry up to 3 times"
4483            );
4484        });
4485
4486        // Advance clock for first retry
4487        cx.executor().advance_clock(BASE_RETRY_DELAY);
4488        cx.run_until_parked();
4489
4490        // Advance clock for second retry
4491        cx.executor().advance_clock(BASE_RETRY_DELAY);
4492        cx.run_until_parked();
4493
4494        // Advance clock for third retry
4495        cx.executor().advance_clock(BASE_RETRY_DELAY);
4496        cx.run_until_parked();
4497
4498        // Should have completed all retries - count retry messages
4499        let retry_count = thread.update(cx, |thread, _| {
4500            thread
4501                .messages
4502                .iter()
4503                .filter(|m| {
4504                    m.ui_only
4505                        && m.segments.iter().any(|s| {
4506                            if let MessageSegment::Text(text) = s {
4507                                text.contains("Retrying") && text.contains("seconds")
4508                            } else {
4509                                false
4510                            }
4511                        })
4512                })
4513                .count()
4514        });
4515        assert_eq!(
4516            retry_count, 3,
4517            "Should have 3 retries for internal server errors"
4518        );
4519
4520        // For internal server errors, we retry 3 times and then give up
4521        // Check that retry_state is cleared after all retries
4522        thread.read_with(cx, |thread, _| {
4523            assert!(
4524                thread.retry_state.is_none(),
4525                "Retry state should be cleared after all retries"
4526            );
4527        });
4528
4529        // Verify total attempts (1 initial + 3 retries)
4530        assert_eq!(
4531            *completion_count.lock(),
4532            4,
4533            "Should have attempted once plus 3 retries"
4534        );
4535    }
4536
4537    #[gpui::test]
4538    async fn test_max_retries_exceeded(cx: &mut TestAppContext) {
4539        init_test_settings(cx);
4540
4541        let project = create_test_project(cx, json!({})).await;
4542        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4543
4544        // Enable Burn Mode to allow retries
4545        thread.update(cx, |thread, _| {
4546            thread.set_completion_mode(CompletionMode::Burn);
4547        });
4548
4549        // Create model that returns overloaded error
4550        let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
4551
4552        // Insert a user message
4553        thread.update(cx, |thread, cx| {
4554            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4555        });
4556
4557        // Track events
4558        let stopped_with_error = Arc::new(Mutex::new(false));
4559        let stopped_with_error_clone = stopped_with_error.clone();
4560
4561        let _subscription = thread.update(cx, |_, cx| {
4562            cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4563                if let ThreadEvent::Stopped(Err(_)) = event {
4564                    *stopped_with_error_clone.lock() = true;
4565                }
4566            })
4567        });
4568
4569        // Start initial completion
4570        thread.update(cx, |thread, cx| {
4571            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4572        });
4573        cx.run_until_parked();
4574
4575        // Advance through all retries
4576        for _ in 0..MAX_RETRY_ATTEMPTS {
4577            cx.executor().advance_clock(BASE_RETRY_DELAY);
4578            cx.run_until_parked();
4579        }
4580
4581        let retry_count = thread.update(cx, |thread, _| {
4582            thread
4583                .messages
4584                .iter()
4585                .filter(|m| {
4586                    m.ui_only
4587                        && m.segments.iter().any(|s| {
4588                            if let MessageSegment::Text(text) = s {
4589                                text.contains("Retrying") && text.contains("seconds")
4590                            } else {
4591                                false
4592                            }
4593                        })
4594                })
4595                .count()
4596        });
4597
4598        // After max retries, should emit Stopped(Err(...)) event
4599        assert_eq!(
4600            retry_count, MAX_RETRY_ATTEMPTS as usize,
4601            "Should have attempted MAX_RETRY_ATTEMPTS retries for overloaded errors"
4602        );
4603        assert!(
4604            *stopped_with_error.lock(),
4605            "Should emit Stopped(Err(...)) event after max retries exceeded"
4606        );
4607
4608        // Retry state should be cleared
4609        thread.read_with(cx, |thread, _| {
4610            assert!(
4611                thread.retry_state.is_none(),
4612                "Retry state should be cleared after max retries"
4613            );
4614
4615            // Verify we have the expected number of retry messages
4616            let retry_messages = thread
4617                .messages
4618                .iter()
4619                .filter(|msg| msg.ui_only && msg.role == Role::System)
4620                .count();
4621            assert_eq!(
4622                retry_messages, MAX_RETRY_ATTEMPTS as usize,
4623                "Should have MAX_RETRY_ATTEMPTS retry messages for overloaded errors"
4624            );
4625        });
4626    }
4627
4628    #[gpui::test]
4629    async fn test_retry_message_removed_on_retry(cx: &mut TestAppContext) {
4630        init_test_settings(cx);
4631
4632        let project = create_test_project(cx, json!({})).await;
4633        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4634
4635        // Enable Burn Mode to allow retries
4636        thread.update(cx, |thread, _| {
4637            thread.set_completion_mode(CompletionMode::Burn);
4638        });
4639
4640        // We'll use a wrapper to switch behavior after first failure
4641        struct RetryTestModel {
4642            inner: Arc<FakeLanguageModel>,
4643            failed_once: Arc<Mutex<bool>>,
4644        }
4645
4646        impl LanguageModel for RetryTestModel {
4647            fn id(&self) -> LanguageModelId {
4648                self.inner.id()
4649            }
4650
4651            fn name(&self) -> LanguageModelName {
4652                self.inner.name()
4653            }
4654
4655            fn provider_id(&self) -> LanguageModelProviderId {
4656                self.inner.provider_id()
4657            }
4658
4659            fn provider_name(&self) -> LanguageModelProviderName {
4660                self.inner.provider_name()
4661            }
4662
4663            fn supports_tools(&self) -> bool {
4664                self.inner.supports_tools()
4665            }
4666
4667            fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4668                self.inner.supports_tool_choice(choice)
4669            }
4670
4671            fn supports_images(&self) -> bool {
4672                self.inner.supports_images()
4673            }
4674
4675            fn telemetry_id(&self) -> String {
4676                self.inner.telemetry_id()
4677            }
4678
4679            fn max_token_count(&self) -> u64 {
4680                self.inner.max_token_count()
4681            }
4682
4683            fn count_tokens(
4684                &self,
4685                request: LanguageModelRequest,
4686                cx: &App,
4687            ) -> BoxFuture<'static, Result<u64>> {
4688                self.inner.count_tokens(request, cx)
4689            }
4690
4691            fn stream_completion(
4692                &self,
4693                request: LanguageModelRequest,
4694                cx: &AsyncApp,
4695            ) -> BoxFuture<
4696                'static,
4697                Result<
4698                    BoxStream<
4699                        'static,
4700                        Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4701                    >,
4702                    LanguageModelCompletionError,
4703                >,
4704            > {
4705                if !*self.failed_once.lock() {
4706                    *self.failed_once.lock() = true;
4707                    let provider = self.provider_name();
4708                    // Return error on first attempt
4709                    let stream = futures::stream::once(async move {
4710                        Err(LanguageModelCompletionError::ServerOverloaded {
4711                            provider,
4712                            retry_after: None,
4713                        })
4714                    });
4715                    async move { Ok(stream.boxed()) }.boxed()
4716                } else {
4717                    // Succeed on retry
4718                    self.inner.stream_completion(request, cx)
4719                }
4720            }
4721
4722            fn as_fake(&self) -> &FakeLanguageModel {
4723                &self.inner
4724            }
4725        }
4726
4727        let model = Arc::new(RetryTestModel {
4728            inner: Arc::new(FakeLanguageModel::default()),
4729            failed_once: Arc::new(Mutex::new(false)),
4730        });
4731
4732        // Insert a user message
4733        thread.update(cx, |thread, cx| {
4734            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
4735        });
4736
4737        // Track message deletions
4738        // Track when retry completes successfully
4739        let retry_completed = Arc::new(Mutex::new(false));
4740        let retry_completed_clone = retry_completed.clone();
4741
4742        let _subscription = thread.update(cx, |_, cx| {
4743            cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
4744                if let ThreadEvent::StreamedCompletion = event {
4745                    *retry_completed_clone.lock() = true;
4746                }
4747            })
4748        });
4749
4750        // Start completion
4751        thread.update(cx, |thread, cx| {
4752            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
4753        });
4754        cx.run_until_parked();
4755
4756        // Get the retry message ID
4757        let retry_message_id = thread.read_with(cx, |thread, _| {
4758            thread
4759                .messages()
4760                .find(|msg| msg.role == Role::System && msg.ui_only)
4761                .map(|msg| msg.id)
4762                .expect("Should have a retry message")
4763        });
4764
4765        // Wait for retry
4766        cx.executor().advance_clock(BASE_RETRY_DELAY);
4767        cx.run_until_parked();
4768
4769        // Stream some successful content
4770        let fake_model = model.as_fake();
4771        // After the retry, there should be a new pending completion
4772        let pending = fake_model.pending_completions();
4773        assert!(
4774            !pending.is_empty(),
4775            "Should have a pending completion after retry"
4776        );
4777        fake_model.stream_completion_response(&pending[0], "Success!");
4778        fake_model.end_completion_stream(&pending[0]);
4779        cx.run_until_parked();
4780
4781        // Check that the retry completed successfully
4782        assert!(
4783            *retry_completed.lock(),
4784            "Retry should have completed successfully"
4785        );
4786
4787        // Retry message should still exist but be marked as ui_only
4788        thread.read_with(cx, |thread, _| {
4789            let retry_msg = thread
4790                .message(retry_message_id)
4791                .expect("Retry message should still exist");
4792            assert!(retry_msg.ui_only, "Retry message should be ui_only");
4793            assert_eq!(
4794                retry_msg.role,
4795                Role::System,
4796                "Retry message should have System role"
4797            );
4798        });
4799    }
4800
4801    #[gpui::test]
4802    async fn test_successful_completion_clears_retry_state(cx: &mut TestAppContext) {
4803        init_test_settings(cx);
4804
4805        let project = create_test_project(cx, json!({})).await;
4806        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4807
4808        // Enable Burn Mode to allow retries
4809        thread.update(cx, |thread, _| {
4810            thread.set_completion_mode(CompletionMode::Burn);
4811        });
4812
4813        // Create a model that fails once then succeeds
4814        struct FailOnceModel {
4815            inner: Arc<FakeLanguageModel>,
4816            failed_once: Arc<Mutex<bool>>,
4817        }
4818
4819        impl LanguageModel for FailOnceModel {
4820            fn id(&self) -> LanguageModelId {
4821                self.inner.id()
4822            }
4823
4824            fn name(&self) -> LanguageModelName {
4825                self.inner.name()
4826            }
4827
4828            fn provider_id(&self) -> LanguageModelProviderId {
4829                self.inner.provider_id()
4830            }
4831
4832            fn provider_name(&self) -> LanguageModelProviderName {
4833                self.inner.provider_name()
4834            }
4835
4836            fn supports_tools(&self) -> bool {
4837                self.inner.supports_tools()
4838            }
4839
4840            fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
4841                self.inner.supports_tool_choice(choice)
4842            }
4843
4844            fn supports_images(&self) -> bool {
4845                self.inner.supports_images()
4846            }
4847
4848            fn telemetry_id(&self) -> String {
4849                self.inner.telemetry_id()
4850            }
4851
4852            fn max_token_count(&self) -> u64 {
4853                self.inner.max_token_count()
4854            }
4855
4856            fn count_tokens(
4857                &self,
4858                request: LanguageModelRequest,
4859                cx: &App,
4860            ) -> BoxFuture<'static, Result<u64>> {
4861                self.inner.count_tokens(request, cx)
4862            }
4863
4864            fn stream_completion(
4865                &self,
4866                request: LanguageModelRequest,
4867                cx: &AsyncApp,
4868            ) -> BoxFuture<
4869                'static,
4870                Result<
4871                    BoxStream<
4872                        'static,
4873                        Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
4874                    >,
4875                    LanguageModelCompletionError,
4876                >,
4877            > {
4878                if !*self.failed_once.lock() {
4879                    *self.failed_once.lock() = true;
4880                    let provider = self.provider_name();
4881                    // Return error on first attempt
4882                    let stream = futures::stream::once(async move {
4883                        Err(LanguageModelCompletionError::ServerOverloaded {
4884                            provider,
4885                            retry_after: None,
4886                        })
4887                    });
4888                    async move { Ok(stream.boxed()) }.boxed()
4889                } else {
4890                    // Succeed on retry
4891                    self.inner.stream_completion(request, cx)
4892                }
4893            }
4894        }
4895
4896        let fail_once_model = Arc::new(FailOnceModel {
4897            inner: Arc::new(FakeLanguageModel::default()),
4898            failed_once: Arc::new(Mutex::new(false)),
4899        });
4900
4901        // Insert a user message
4902        thread.update(cx, |thread, cx| {
4903            thread.insert_user_message(
4904                "Test message",
4905                ContextLoadResult::default(),
4906                None,
4907                vec![],
4908                cx,
4909            );
4910        });
4911
4912        // Start completion with fail-once model
4913        thread.update(cx, |thread, cx| {
4914            thread.send_to_model(
4915                fail_once_model.clone(),
4916                CompletionIntent::UserPrompt,
4917                None,
4918                cx,
4919            );
4920        });
4921
4922        cx.run_until_parked();
4923
4924        // Verify retry state exists after first failure
4925        thread.read_with(cx, |thread, _| {
4926            assert!(
4927                thread.retry_state.is_some(),
4928                "Should have retry state after failure"
4929            );
4930        });
4931
4932        // Wait for retry delay
4933        cx.executor().advance_clock(BASE_RETRY_DELAY);
4934        cx.run_until_parked();
4935
4936        // The retry should now use our FailOnceModel which should succeed
4937        // We need to help the FakeLanguageModel complete the stream
4938        let inner_fake = fail_once_model.inner.clone();
4939
4940        // Wait a bit for the retry to start
4941        cx.run_until_parked();
4942
4943        // Check for pending completions and complete them
4944        if let Some(pending) = inner_fake.pending_completions().first() {
4945            inner_fake.stream_completion_response(pending, "Success!");
4946            inner_fake.end_completion_stream(pending);
4947        }
4948        cx.run_until_parked();
4949
4950        thread.read_with(cx, |thread, _| {
4951            assert!(
4952                thread.retry_state.is_none(),
4953                "Retry state should be cleared after successful completion"
4954            );
4955
4956            let has_assistant_message = thread
4957                .messages
4958                .iter()
4959                .any(|msg| msg.role == Role::Assistant && !msg.ui_only);
4960            assert!(
4961                has_assistant_message,
4962                "Should have an assistant message after successful retry"
4963            );
4964        });
4965    }
4966
4967    #[gpui::test]
4968    async fn test_rate_limit_retry_single_attempt(cx: &mut TestAppContext) {
4969        init_test_settings(cx);
4970
4971        let project = create_test_project(cx, json!({})).await;
4972        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
4973
4974        // Enable Burn Mode to allow retries
4975        thread.update(cx, |thread, _| {
4976            thread.set_completion_mode(CompletionMode::Burn);
4977        });
4978
4979        // Create a model that returns rate limit error with retry_after
4980        struct RateLimitModel {
4981            inner: Arc<FakeLanguageModel>,
4982        }
4983
4984        impl LanguageModel for RateLimitModel {
4985            fn id(&self) -> LanguageModelId {
4986                self.inner.id()
4987            }
4988
4989            fn name(&self) -> LanguageModelName {
4990                self.inner.name()
4991            }
4992
4993            fn provider_id(&self) -> LanguageModelProviderId {
4994                self.inner.provider_id()
4995            }
4996
4997            fn provider_name(&self) -> LanguageModelProviderName {
4998                self.inner.provider_name()
4999            }
5000
5001            fn supports_tools(&self) -> bool {
5002                self.inner.supports_tools()
5003            }
5004
5005            fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
5006                self.inner.supports_tool_choice(choice)
5007            }
5008
5009            fn supports_images(&self) -> bool {
5010                self.inner.supports_images()
5011            }
5012
5013            fn telemetry_id(&self) -> String {
5014                self.inner.telemetry_id()
5015            }
5016
5017            fn max_token_count(&self) -> u64 {
5018                self.inner.max_token_count()
5019            }
5020
5021            fn count_tokens(
5022                &self,
5023                request: LanguageModelRequest,
5024                cx: &App,
5025            ) -> BoxFuture<'static, Result<u64>> {
5026                self.inner.count_tokens(request, cx)
5027            }
5028
5029            fn stream_completion(
5030                &self,
5031                _request: LanguageModelRequest,
5032                _cx: &AsyncApp,
5033            ) -> BoxFuture<
5034                'static,
5035                Result<
5036                    BoxStream<
5037                        'static,
5038                        Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
5039                    >,
5040                    LanguageModelCompletionError,
5041                >,
5042            > {
5043                let provider = self.provider_name();
5044                async move {
5045                    let stream = futures::stream::once(async move {
5046                        Err(LanguageModelCompletionError::RateLimitExceeded {
5047                            provider,
5048                            retry_after: Some(Duration::from_secs(TEST_RATE_LIMIT_RETRY_SECS)),
5049                        })
5050                    });
5051                    Ok(stream.boxed())
5052                }
5053                .boxed()
5054            }
5055
5056            fn as_fake(&self) -> &FakeLanguageModel {
5057                &self.inner
5058            }
5059        }
5060
5061        let model = Arc::new(RateLimitModel {
5062            inner: Arc::new(FakeLanguageModel::default()),
5063        });
5064
5065        // Insert a user message
5066        thread.update(cx, |thread, cx| {
5067            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5068        });
5069
5070        // Start completion
5071        thread.update(cx, |thread, cx| {
5072            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5073        });
5074
5075        cx.run_until_parked();
5076
5077        let retry_count = thread.update(cx, |thread, _| {
5078            thread
5079                .messages
5080                .iter()
5081                .filter(|m| {
5082                    m.ui_only
5083                        && m.segments.iter().any(|s| {
5084                            if let MessageSegment::Text(text) = s {
5085                                text.contains("rate limit exceeded")
5086                            } else {
5087                                false
5088                            }
5089                        })
5090                })
5091                .count()
5092        });
5093        assert_eq!(retry_count, 1, "Should have scheduled one retry");
5094
5095        thread.read_with(cx, |thread, _| {
5096            assert!(
5097                thread.retry_state.is_some(),
5098                "Rate limit errors should set retry_state"
5099            );
5100            if let Some(retry_state) = &thread.retry_state {
5101                assert_eq!(
5102                    retry_state.max_attempts, MAX_RETRY_ATTEMPTS,
5103                    "Rate limit errors should use MAX_RETRY_ATTEMPTS"
5104                );
5105            }
5106        });
5107
5108        // Verify we have one retry message
5109        thread.read_with(cx, |thread, _| {
5110            let retry_messages = thread
5111                .messages
5112                .iter()
5113                .filter(|msg| {
5114                    msg.ui_only
5115                        && msg.segments.iter().any(|seg| {
5116                            if let MessageSegment::Text(text) = seg {
5117                                text.contains("rate limit exceeded")
5118                            } else {
5119                                false
5120                            }
5121                        })
5122                })
5123                .count();
5124            assert_eq!(
5125                retry_messages, 1,
5126                "Should have one rate limit retry message"
5127            );
5128        });
5129
5130        // Check that retry message doesn't include attempt count
5131        thread.read_with(cx, |thread, _| {
5132            let retry_message = thread
5133                .messages
5134                .iter()
5135                .find(|msg| msg.role == Role::System && msg.ui_only)
5136                .expect("Should have a retry message");
5137
5138            // Check that the message contains attempt count since we use retry_state
5139            if let Some(MessageSegment::Text(text)) = retry_message.segments.first() {
5140                assert!(
5141                    text.contains(&format!("attempt 1 of {}", MAX_RETRY_ATTEMPTS)),
5142                    "Rate limit retry message should contain attempt count with MAX_RETRY_ATTEMPTS"
5143                );
5144                assert!(
5145                    text.contains("Retrying"),
5146                    "Rate limit retry message should contain retry text"
5147                );
5148            }
5149        });
5150    }
5151
5152    #[gpui::test]
5153    async fn test_ui_only_messages_not_sent_to_model(cx: &mut TestAppContext) {
5154        init_test_settings(cx);
5155
5156        let project = create_test_project(cx, json!({})).await;
5157        let (_, _, thread, _, model) = setup_test_environment(cx, project.clone()).await;
5158
5159        // Insert a regular user message
5160        thread.update(cx, |thread, cx| {
5161            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5162        });
5163
5164        // Insert a UI-only message (like our retry notifications)
5165        thread.update(cx, |thread, cx| {
5166            let id = thread.next_message_id.post_inc();
5167            thread.messages.push(Message {
5168                id,
5169                role: Role::System,
5170                segments: vec![MessageSegment::Text(
5171                    "This is a UI-only message that should not be sent to the model".to_string(),
5172                )],
5173                loaded_context: LoadedContext::default(),
5174                creases: Vec::new(),
5175                is_hidden: true,
5176                ui_only: true,
5177            });
5178            cx.emit(ThreadEvent::MessageAdded(id));
5179        });
5180
5181        // Insert another regular message
5182        thread.update(cx, |thread, cx| {
5183            thread.insert_user_message(
5184                "How are you?",
5185                ContextLoadResult::default(),
5186                None,
5187                vec![],
5188                cx,
5189            );
5190        });
5191
5192        // Generate the completion request
5193        let request = thread.update(cx, |thread, cx| {
5194            thread.to_completion_request(model.clone(), CompletionIntent::UserPrompt, cx)
5195        });
5196
5197        // Verify that the request only contains non-UI-only messages
5198        // Should have system prompt + 2 user messages, but not the UI-only message
5199        let user_messages: Vec<_> = request
5200            .messages
5201            .iter()
5202            .filter(|msg| msg.role == Role::User)
5203            .collect();
5204        assert_eq!(
5205            user_messages.len(),
5206            2,
5207            "Should have exactly 2 user messages"
5208        );
5209
5210        // Verify the UI-only content is not present anywhere in the request
5211        let request_text = request
5212            .messages
5213            .iter()
5214            .flat_map(|msg| &msg.content)
5215            .filter_map(|content| match content {
5216                MessageContent::Text(text) => Some(text.as_str()),
5217                _ => None,
5218            })
5219            .collect::<String>();
5220
5221        assert!(
5222            !request_text.contains("UI-only message"),
5223            "UI-only message content should not be in the request"
5224        );
5225
5226        // Verify the thread still has all 3 messages (including UI-only)
5227        thread.read_with(cx, |thread, _| {
5228            assert_eq!(
5229                thread.messages().count(),
5230                3,
5231                "Thread should have 3 messages"
5232            );
5233            assert_eq!(
5234                thread.messages().filter(|m| m.ui_only).count(),
5235                1,
5236                "Thread should have 1 UI-only message"
5237            );
5238        });
5239
5240        // Verify that UI-only messages are not serialized
5241        let serialized = thread
5242            .update(cx, |thread, cx| thread.serialize(cx))
5243            .await
5244            .unwrap();
5245        assert_eq!(
5246            serialized.messages.len(),
5247            2,
5248            "Serialized thread should only have 2 messages (no UI-only)"
5249        );
5250    }
5251
5252    #[gpui::test]
5253    async fn test_no_retry_without_burn_mode(cx: &mut TestAppContext) {
5254        init_test_settings(cx);
5255
5256        let project = create_test_project(cx, json!({})).await;
5257        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
5258
5259        // Ensure we're in Normal mode (not Burn mode)
5260        thread.update(cx, |thread, _| {
5261            thread.set_completion_mode(CompletionMode::Normal);
5262        });
5263
5264        // Track error events
5265        let error_events = Arc::new(Mutex::new(Vec::new()));
5266        let error_events_clone = error_events.clone();
5267
5268        let _subscription = thread.update(cx, |_, cx| {
5269            cx.subscribe(&thread, move |_, _, event: &ThreadEvent, _| {
5270                if let ThreadEvent::ShowError(error) = event {
5271                    error_events_clone.lock().push(error.clone());
5272                }
5273            })
5274        });
5275
5276        // Create model that returns overloaded error
5277        let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
5278
5279        // Insert a user message
5280        thread.update(cx, |thread, cx| {
5281            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5282        });
5283
5284        // Start completion
5285        thread.update(cx, |thread, cx| {
5286            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5287        });
5288
5289        cx.run_until_parked();
5290
5291        // Verify no retry state was created
5292        thread.read_with(cx, |thread, _| {
5293            assert!(
5294                thread.retry_state.is_none(),
5295                "Should not have retry state in Normal mode"
5296            );
5297        });
5298
5299        // Check that a retryable error was reported
5300        let errors = error_events.lock();
5301        assert!(!errors.is_empty(), "Should have received an error event");
5302
5303        if let ThreadError::RetryableError {
5304            message: _,
5305            can_enable_burn_mode,
5306        } = &errors[0]
5307        {
5308            assert!(
5309                *can_enable_burn_mode,
5310                "Error should indicate burn mode can be enabled"
5311            );
5312        } else {
5313            panic!("Expected RetryableError, got {:?}", errors[0]);
5314        }
5315
5316        // Verify the thread is no longer generating
5317        thread.read_with(cx, |thread, _| {
5318            assert!(
5319                !thread.is_generating(),
5320                "Should not be generating after error without retry"
5321            );
5322        });
5323    }
5324
5325    #[gpui::test]
5326    async fn test_retry_cancelled_on_stop(cx: &mut TestAppContext) {
5327        init_test_settings(cx);
5328
5329        let project = create_test_project(cx, json!({})).await;
5330        let (_, _, thread, _, _base_model) = setup_test_environment(cx, project.clone()).await;
5331
5332        // Enable Burn Mode to allow retries
5333        thread.update(cx, |thread, _| {
5334            thread.set_completion_mode(CompletionMode::Burn);
5335        });
5336
5337        // Create model that returns overloaded error
5338        let model = Arc::new(ErrorInjector::new(TestError::Overloaded));
5339
5340        // Insert a user message
5341        thread.update(cx, |thread, cx| {
5342            thread.insert_user_message("Hello!", ContextLoadResult::default(), None, vec![], cx);
5343        });
5344
5345        // Start completion
5346        thread.update(cx, |thread, cx| {
5347            thread.send_to_model(model.clone(), CompletionIntent::UserPrompt, None, cx);
5348        });
5349
5350        cx.run_until_parked();
5351
5352        // Verify retry was scheduled by checking for retry message
5353        let has_retry_message = thread.read_with(cx, |thread, _| {
5354            thread.messages.iter().any(|m| {
5355                m.ui_only
5356                    && m.segments.iter().any(|s| {
5357                        if let MessageSegment::Text(text) = s {
5358                            text.contains("Retrying") && text.contains("seconds")
5359                        } else {
5360                            false
5361                        }
5362                    })
5363            })
5364        });
5365        assert!(has_retry_message, "Should have scheduled a retry");
5366
5367        // Cancel the completion before the retry happens
5368        thread.update(cx, |thread, cx| {
5369            thread.cancel_last_completion(None, cx);
5370        });
5371
5372        cx.run_until_parked();
5373
5374        // The retry should not have happened - no pending completions
5375        let fake_model = model.as_fake();
5376        assert_eq!(
5377            fake_model.pending_completions().len(),
5378            0,
5379            "Should have no pending completions after cancellation"
5380        );
5381
5382        // Verify the retry was cancelled by checking retry state
5383        thread.read_with(cx, |thread, _| {
5384            if let Some(retry_state) = &thread.retry_state {
5385                panic!(
5386                    "retry_state should be cleared after cancellation, but found: attempt={}, max_attempts={}, intent={:?}",
5387                    retry_state.attempt, retry_state.max_attempts, retry_state.intent
5388                );
5389            }
5390        });
5391    }
5392
5393    fn test_summarize_error(
5394        model: &Arc<dyn LanguageModel>,
5395        thread: &Entity<Thread>,
5396        cx: &mut TestAppContext,
5397    ) {
5398        thread.update(cx, |thread, cx| {
5399            thread.insert_user_message("Hi!", ContextLoadResult::default(), None, vec![], cx);
5400            thread.send_to_model(
5401                model.clone(),
5402                CompletionIntent::ThreadSummarization,
5403                None,
5404                cx,
5405            );
5406        });
5407
5408        let fake_model = model.as_fake();
5409        simulate_successful_response(&fake_model, cx);
5410
5411        thread.read_with(cx, |thread, _| {
5412            assert!(matches!(thread.summary(), ThreadSummary::Generating));
5413            assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
5414        });
5415
5416        // Simulate summary request ending
5417        cx.run_until_parked();
5418        fake_model.end_last_completion_stream();
5419        cx.run_until_parked();
5420
5421        // State is set to Error and default message
5422        thread.read_with(cx, |thread, _| {
5423            assert!(matches!(thread.summary(), ThreadSummary::Error));
5424            assert_eq!(thread.summary().or_default(), ThreadSummary::DEFAULT);
5425        });
5426    }
5427
5428    fn simulate_successful_response(fake_model: &FakeLanguageModel, cx: &mut TestAppContext) {
5429        cx.run_until_parked();
5430        fake_model.stream_last_completion_response("Assistant response");
5431        fake_model.end_last_completion_stream();
5432        cx.run_until_parked();
5433    }
5434
5435    fn init_test_settings(cx: &mut TestAppContext) {
5436        cx.update(|cx| {
5437            let settings_store = SettingsStore::test(cx);
5438            cx.set_global(settings_store);
5439            language::init(cx);
5440            Project::init_settings(cx);
5441            AgentSettings::register(cx);
5442            prompt_store::init(cx);
5443            thread_store::init(cx);
5444            workspace::init_settings(cx);
5445            language_model::init_settings(cx);
5446            ThemeSettings::register(cx);
5447            ToolRegistry::default_global(cx);
5448            assistant_tool::init(cx);
5449
5450            let http_client = Arc::new(http_client::HttpClientWithUrl::new(
5451                http_client::FakeHttpClient::with_200_response(),
5452                "http://localhost".to_string(),
5453                None,
5454            ));
5455            assistant_tools::init(http_client, cx);
5456        });
5457    }
5458
5459    // Helper to create a test project with test files
5460    async fn create_test_project(
5461        cx: &mut TestAppContext,
5462        files: serde_json::Value,
5463    ) -> Entity<Project> {
5464        let fs = FakeFs::new(cx.executor());
5465        fs.insert_tree(path!("/test"), files).await;
5466        Project::test(fs, [path!("/test").as_ref()], cx).await
5467    }
5468
5469    async fn setup_test_environment(
5470        cx: &mut TestAppContext,
5471        project: Entity<Project>,
5472    ) -> (
5473        Entity<Workspace>,
5474        Entity<ThreadStore>,
5475        Entity<Thread>,
5476        Entity<ContextStore>,
5477        Arc<dyn LanguageModel>,
5478    ) {
5479        let (workspace, cx) =
5480            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
5481
5482        let thread_store = cx
5483            .update(|_, cx| {
5484                ThreadStore::load(
5485                    project.clone(),
5486                    cx.new(|_| ToolWorkingSet::default()),
5487                    None,
5488                    Arc::new(PromptBuilder::new(None).unwrap()),
5489                    cx,
5490                )
5491            })
5492            .await
5493            .unwrap();
5494
5495        let thread = thread_store.update(cx, |store, cx| store.create_thread(cx));
5496        let context_store = cx.new(|_cx| ContextStore::new(project.downgrade(), None));
5497
5498        let provider = Arc::new(FakeLanguageModelProvider::default());
5499        let model = provider.test_model();
5500        let model: Arc<dyn LanguageModel> = Arc::new(model);
5501
5502        cx.update(|_, cx| {
5503            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
5504                registry.set_default_model(
5505                    Some(ConfiguredModel {
5506                        provider: provider.clone(),
5507                        model: model.clone(),
5508                    }),
5509                    cx,
5510                );
5511                registry.set_thread_summary_model(
5512                    Some(ConfiguredModel {
5513                        provider,
5514                        model: model.clone(),
5515                    }),
5516                    cx,
5517                );
5518            })
5519        });
5520
5521        (workspace, thread_store, thread, context_store, model)
5522    }
5523
5524    async fn add_file_to_context(
5525        project: &Entity<Project>,
5526        context_store: &Entity<ContextStore>,
5527        path: &str,
5528        cx: &mut TestAppContext,
5529    ) -> Result<Entity<language::Buffer>> {
5530        let buffer_path = project
5531            .read_with(cx, |project, cx| project.find_project_path(path, cx))
5532            .unwrap();
5533
5534        let buffer = project
5535            .update(cx, |project, cx| {
5536                project.open_buffer(buffer_path.clone(), cx)
5537            })
5538            .await
5539            .unwrap();
5540
5541        context_store.update(cx, |context_store, cx| {
5542            context_store.add_file_from_buffer(&buffer_path, buffer.clone(), false, cx);
5543        });
5544
5545        Ok(buffer)
5546    }
5547}