thread.rs

   1use crate::{
   2    ContextServerRegistry, CopyPathTool, CreateDirectoryTool, DbLanguageModel, DbThread,
   3    DeletePathTool, DiagnosticsTool, EditFileTool, FetchTool, FindPathTool, GrepTool,
   4    ListDirectoryTool, MovePathTool, NowTool, OpenTool, ReadFileTool, SystemPromptTemplate,
   5    Template, Templates, TerminalTool, ThinkingTool, WebSearchTool,
   6};
   7use acp_thread::{MentionUri, UserMessageId};
   8use action_log::ActionLog;
   9use agent::thread::{GitState, ProjectSnapshot, WorktreeSnapshot};
  10use agent_client_protocol as acp;
  11use agent_settings::{
  12    AgentProfileId, AgentProfileSettings, AgentSettings, CompletionMode,
  13    SUMMARIZE_THREAD_DETAILED_PROMPT, SUMMARIZE_THREAD_PROMPT,
  14};
  15use anyhow::{Context as _, Result, anyhow};
  16use assistant_tool::adapt_schema_to_format;
  17use chrono::{DateTime, Utc};
  18use client::{ModelRequestUsage, RequestUsage, UserStore};
  19use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, Plan, UsageLimit};
  20use collections::{HashMap, HashSet, IndexMap};
  21use fs::Fs;
  22use futures::stream;
  23use futures::{
  24    FutureExt,
  25    channel::{mpsc, oneshot},
  26    future::Shared,
  27    stream::FuturesUnordered,
  28};
  29use git::repository::DiffType;
  30use gpui::{
  31    App, AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity,
  32};
  33use language_model::{
  34    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelExt,
  35    LanguageModelImage, LanguageModelProviderId, LanguageModelRegistry, LanguageModelRequest,
  36    LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
  37    LanguageModelToolResultContent, LanguageModelToolSchemaFormat, LanguageModelToolUse,
  38    LanguageModelToolUseId, Role, SelectedModel, StopReason, TokenUsage, ZED_CLOUD_PROVIDER_ID,
  39};
  40use project::{
  41    Project,
  42    git_store::{GitStore, RepositoryState},
  43};
  44use prompt_store::ProjectContext;
  45use schemars::{JsonSchema, Schema};
  46use serde::{Deserialize, Serialize};
  47use settings::{Settings, update_settings_file};
  48use smol::stream::StreamExt;
  49use std::{
  50    collections::BTreeMap,
  51    ops::RangeInclusive,
  52    path::Path,
  53    rc::Rc,
  54    sync::Arc,
  55    time::{Duration, Instant},
  56};
  57use std::{fmt::Write, path::PathBuf};
  58use util::{ResultExt, debug_panic, markdown::MarkdownCodeBlock};
  59use uuid::Uuid;
  60
  61const TOOL_CANCELED_MESSAGE: &str = "Tool canceled by user";
  62pub const MAX_TOOL_NAME_LENGTH: usize = 64;
  63
  64/// The ID of the user prompt that initiated a request.
  65///
  66/// This equates to the user physically submitting a message to the model (e.g., by pressing the Enter key).
  67#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
  68pub struct PromptId(Arc<str>);
  69
  70impl PromptId {
  71    pub fn new() -> Self {
  72        Self(Uuid::new_v4().to_string().into())
  73    }
  74}
  75
  76impl std::fmt::Display for PromptId {
  77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
  78        write!(f, "{}", self.0)
  79    }
  80}
  81
  82pub(crate) const MAX_RETRY_ATTEMPTS: u8 = 4;
  83pub(crate) const BASE_RETRY_DELAY: Duration = Duration::from_secs(5);
  84
  85#[derive(Debug, Clone)]
  86enum RetryStrategy {
  87    ExponentialBackoff {
  88        initial_delay: Duration,
  89        max_attempts: u8,
  90    },
  91    Fixed {
  92        delay: Duration,
  93        max_attempts: u8,
  94    },
  95}
  96
  97#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
  98pub enum Message {
  99    User(UserMessage),
 100    Agent(AgentMessage),
 101    Resume,
 102}
 103
 104impl Message {
 105    pub fn as_agent_message(&self) -> Option<&AgentMessage> {
 106        match self {
 107            Message::Agent(agent_message) => Some(agent_message),
 108            _ => None,
 109        }
 110    }
 111
 112    pub fn to_request(&self) -> Vec<LanguageModelRequestMessage> {
 113        match self {
 114            Message::User(message) => vec![message.to_request()],
 115            Message::Agent(message) => message.to_request(),
 116            Message::Resume => vec![LanguageModelRequestMessage {
 117                role: Role::User,
 118                content: vec!["Continue where you left off".into()],
 119                cache: false,
 120            }],
 121        }
 122    }
 123
 124    pub fn to_markdown(&self) -> String {
 125        match self {
 126            Message::User(message) => message.to_markdown(),
 127            Message::Agent(message) => message.to_markdown(),
 128            Message::Resume => "[resume]\n".into(),
 129        }
 130    }
 131
 132    pub fn role(&self) -> Role {
 133        match self {
 134            Message::User(_) | Message::Resume => Role::User,
 135            Message::Agent(_) => Role::Assistant,
 136        }
 137    }
 138}
 139
 140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 141pub struct UserMessage {
 142    pub id: UserMessageId,
 143    pub content: Vec<UserMessageContent>,
 144}
 145
 146#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 147pub enum UserMessageContent {
 148    Text(String),
 149    Mention { uri: MentionUri, content: String },
 150    Image(LanguageModelImage),
 151}
 152
 153impl UserMessage {
 154    pub fn to_markdown(&self) -> String {
 155        let mut markdown = String::from("## User\n\n");
 156
 157        for content in &self.content {
 158            match content {
 159                UserMessageContent::Text(text) => {
 160                    markdown.push_str(text);
 161                    markdown.push('\n');
 162                }
 163                UserMessageContent::Image(_) => {
 164                    markdown.push_str("<image />\n");
 165                }
 166                UserMessageContent::Mention { uri, content } => {
 167                    if !content.is_empty() {
 168                        let _ = writeln!(&mut markdown, "{}\n\n{}", uri.as_link(), content);
 169                    } else {
 170                        let _ = writeln!(&mut markdown, "{}", uri.as_link());
 171                    }
 172                }
 173            }
 174        }
 175
 176        markdown
 177    }
 178
 179    fn to_request(&self) -> LanguageModelRequestMessage {
 180        let mut message = LanguageModelRequestMessage {
 181            role: Role::User,
 182            content: Vec::with_capacity(self.content.len()),
 183            cache: false,
 184        };
 185
 186        const OPEN_CONTEXT: &str = "<context>\n\
 187            The following items were attached by the user. \
 188            They are up-to-date and don't need to be re-read.\n\n";
 189
 190        const OPEN_FILES_TAG: &str = "<files>";
 191        const OPEN_DIRECTORIES_TAG: &str = "<directories>";
 192        const OPEN_SYMBOLS_TAG: &str = "<symbols>";
 193        const OPEN_SELECTIONS_TAG: &str = "<selections>";
 194        const OPEN_THREADS_TAG: &str = "<threads>";
 195        const OPEN_FETCH_TAG: &str = "<fetched_urls>";
 196        const OPEN_RULES_TAG: &str =
 197            "<rules>\nThe user has specified the following rules that should be applied:\n";
 198
 199        let mut file_context = OPEN_FILES_TAG.to_string();
 200        let mut directory_context = OPEN_DIRECTORIES_TAG.to_string();
 201        let mut symbol_context = OPEN_SYMBOLS_TAG.to_string();
 202        let mut selection_context = OPEN_SELECTIONS_TAG.to_string();
 203        let mut thread_context = OPEN_THREADS_TAG.to_string();
 204        let mut fetch_context = OPEN_FETCH_TAG.to_string();
 205        let mut rules_context = OPEN_RULES_TAG.to_string();
 206
 207        for chunk in &self.content {
 208            let chunk = match chunk {
 209                UserMessageContent::Text(text) => {
 210                    language_model::MessageContent::Text(text.clone())
 211                }
 212                UserMessageContent::Image(value) => {
 213                    language_model::MessageContent::Image(value.clone())
 214                }
 215                UserMessageContent::Mention { uri, content } => {
 216                    match uri {
 217                        MentionUri::File { abs_path } => {
 218                            write!(
 219                                &mut file_context,
 220                                "\n{}",
 221                                MarkdownCodeBlock {
 222                                    tag: &codeblock_tag(abs_path, None),
 223                                    text: &content.to_string(),
 224                                }
 225                            )
 226                            .ok();
 227                        }
 228                        MentionUri::PastedImage => {
 229                            debug_panic!("pasted image URI should not be used in mention content")
 230                        }
 231                        MentionUri::Directory { .. } => {
 232                            write!(&mut directory_context, "\n{}\n", content).ok();
 233                        }
 234                        MentionUri::Symbol {
 235                            abs_path: path,
 236                            line_range,
 237                            ..
 238                        } => {
 239                            write!(
 240                                &mut symbol_context,
 241                                "\n{}",
 242                                MarkdownCodeBlock {
 243                                    tag: &codeblock_tag(path, Some(line_range)),
 244                                    text: content
 245                                }
 246                            )
 247                            .ok();
 248                        }
 249                        MentionUri::Selection {
 250                            abs_path: path,
 251                            line_range,
 252                            ..
 253                        } => {
 254                            write!(
 255                                &mut selection_context,
 256                                "\n{}",
 257                                MarkdownCodeBlock {
 258                                    tag: &codeblock_tag(
 259                                        path.as_deref().unwrap_or("Untitled".as_ref()),
 260                                        Some(line_range)
 261                                    ),
 262                                    text: content
 263                                }
 264                            )
 265                            .ok();
 266                        }
 267                        MentionUri::Thread { .. } => {
 268                            write!(&mut thread_context, "\n{}\n", content).ok();
 269                        }
 270                        MentionUri::TextThread { .. } => {
 271                            write!(&mut thread_context, "\n{}\n", content).ok();
 272                        }
 273                        MentionUri::Rule { .. } => {
 274                            write!(
 275                                &mut rules_context,
 276                                "\n{}",
 277                                MarkdownCodeBlock {
 278                                    tag: "",
 279                                    text: content
 280                                }
 281                            )
 282                            .ok();
 283                        }
 284                        MentionUri::Fetch { url } => {
 285                            write!(&mut fetch_context, "\nFetch: {}\n\n{}", url, content).ok();
 286                        }
 287                    }
 288
 289                    language_model::MessageContent::Text(uri.as_link().to_string())
 290                }
 291            };
 292
 293            message.content.push(chunk);
 294        }
 295
 296        let len_before_context = message.content.len();
 297
 298        if file_context.len() > OPEN_FILES_TAG.len() {
 299            file_context.push_str("</files>\n");
 300            message
 301                .content
 302                .push(language_model::MessageContent::Text(file_context));
 303        }
 304
 305        if directory_context.len() > OPEN_DIRECTORIES_TAG.len() {
 306            directory_context.push_str("</directories>\n");
 307            message
 308                .content
 309                .push(language_model::MessageContent::Text(directory_context));
 310        }
 311
 312        if symbol_context.len() > OPEN_SYMBOLS_TAG.len() {
 313            symbol_context.push_str("</symbols>\n");
 314            message
 315                .content
 316                .push(language_model::MessageContent::Text(symbol_context));
 317        }
 318
 319        if selection_context.len() > OPEN_SELECTIONS_TAG.len() {
 320            selection_context.push_str("</selections>\n");
 321            message
 322                .content
 323                .push(language_model::MessageContent::Text(selection_context));
 324        }
 325
 326        if thread_context.len() > OPEN_THREADS_TAG.len() {
 327            thread_context.push_str("</threads>\n");
 328            message
 329                .content
 330                .push(language_model::MessageContent::Text(thread_context));
 331        }
 332
 333        if fetch_context.len() > OPEN_FETCH_TAG.len() {
 334            fetch_context.push_str("</fetched_urls>\n");
 335            message
 336                .content
 337                .push(language_model::MessageContent::Text(fetch_context));
 338        }
 339
 340        if rules_context.len() > OPEN_RULES_TAG.len() {
 341            rules_context.push_str("</user_rules>\n");
 342            message
 343                .content
 344                .push(language_model::MessageContent::Text(rules_context));
 345        }
 346
 347        if message.content.len() > len_before_context {
 348            message.content.insert(
 349                len_before_context,
 350                language_model::MessageContent::Text(OPEN_CONTEXT.into()),
 351            );
 352            message
 353                .content
 354                .push(language_model::MessageContent::Text("</context>".into()));
 355        }
 356
 357        message
 358    }
 359}
 360
 361fn codeblock_tag(full_path: &Path, line_range: Option<&RangeInclusive<u32>>) -> String {
 362    let mut result = String::new();
 363
 364    if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) {
 365        let _ = write!(result, "{} ", extension);
 366    }
 367
 368    let _ = write!(result, "{}", full_path.display());
 369
 370    if let Some(range) = line_range {
 371        if range.start() == range.end() {
 372            let _ = write!(result, ":{}", range.start() + 1);
 373        } else {
 374            let _ = write!(result, ":{}-{}", range.start() + 1, range.end() + 1);
 375        }
 376    }
 377
 378    result
 379}
 380
 381impl AgentMessage {
 382    pub fn to_markdown(&self) -> String {
 383        let mut markdown = String::from("## Assistant\n\n");
 384
 385        for content in &self.content {
 386            match content {
 387                AgentMessageContent::Text(text) => {
 388                    markdown.push_str(text);
 389                    markdown.push('\n');
 390                }
 391                AgentMessageContent::Thinking { text, .. } => {
 392                    markdown.push_str("<think>");
 393                    markdown.push_str(text);
 394                    markdown.push_str("</think>\n");
 395                }
 396                AgentMessageContent::RedactedThinking(_) => {
 397                    markdown.push_str("<redacted_thinking />\n")
 398                }
 399                AgentMessageContent::ToolUse(tool_use) => {
 400                    markdown.push_str(&format!(
 401                        "**Tool Use**: {} (ID: {})\n",
 402                        tool_use.name, tool_use.id
 403                    ));
 404                    markdown.push_str(&format!(
 405                        "{}\n",
 406                        MarkdownCodeBlock {
 407                            tag: "json",
 408                            text: &format!("{:#}", tool_use.input)
 409                        }
 410                    ));
 411                }
 412            }
 413        }
 414
 415        for tool_result in self.tool_results.values() {
 416            markdown.push_str(&format!(
 417                "**Tool Result**: {} (ID: {})\n\n",
 418                tool_result.tool_name, tool_result.tool_use_id
 419            ));
 420            if tool_result.is_error {
 421                markdown.push_str("**ERROR:**\n");
 422            }
 423
 424            match &tool_result.content {
 425                LanguageModelToolResultContent::Text(text) => {
 426                    writeln!(markdown, "{text}\n").ok();
 427                }
 428                LanguageModelToolResultContent::Image(_) => {
 429                    writeln!(markdown, "<image />\n").ok();
 430                }
 431            }
 432
 433            if let Some(output) = tool_result.output.as_ref() {
 434                writeln!(
 435                    markdown,
 436                    "**Debug Output**:\n\n```json\n{}\n```\n",
 437                    serde_json::to_string_pretty(output).unwrap()
 438                )
 439                .unwrap();
 440            }
 441        }
 442
 443        markdown
 444    }
 445
 446    pub fn to_request(&self) -> Vec<LanguageModelRequestMessage> {
 447        let mut assistant_message = LanguageModelRequestMessage {
 448            role: Role::Assistant,
 449            content: Vec::with_capacity(self.content.len()),
 450            cache: false,
 451        };
 452        for chunk in &self.content {
 453            match chunk {
 454                AgentMessageContent::Text(text) => {
 455                    assistant_message
 456                        .content
 457                        .push(language_model::MessageContent::Text(text.clone()));
 458                }
 459                AgentMessageContent::Thinking { text, signature } => {
 460                    assistant_message
 461                        .content
 462                        .push(language_model::MessageContent::Thinking {
 463                            text: text.clone(),
 464                            signature: signature.clone(),
 465                        });
 466                }
 467                AgentMessageContent::RedactedThinking(value) => {
 468                    assistant_message.content.push(
 469                        language_model::MessageContent::RedactedThinking(value.clone()),
 470                    );
 471                }
 472                AgentMessageContent::ToolUse(tool_use) => {
 473                    if self.tool_results.contains_key(&tool_use.id) {
 474                        assistant_message
 475                            .content
 476                            .push(language_model::MessageContent::ToolUse(tool_use.clone()));
 477                    }
 478                }
 479            };
 480        }
 481
 482        let mut user_message = LanguageModelRequestMessage {
 483            role: Role::User,
 484            content: Vec::new(),
 485            cache: false,
 486        };
 487
 488        for tool_result in self.tool_results.values() {
 489            let mut tool_result = tool_result.clone();
 490            // Surprisingly, the API fails if we return an empty string here.
 491            // It thinks we are sending a tool use without a tool result.
 492            if tool_result.content.is_empty() {
 493                tool_result.content = "<Tool returned an empty string>".into();
 494            }
 495            user_message
 496                .content
 497                .push(language_model::MessageContent::ToolResult(tool_result));
 498        }
 499
 500        let mut messages = Vec::new();
 501        if !assistant_message.content.is_empty() {
 502            messages.push(assistant_message);
 503        }
 504        if !user_message.content.is_empty() {
 505            messages.push(user_message);
 506        }
 507        messages
 508    }
 509}
 510
 511#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 512pub struct AgentMessage {
 513    pub content: Vec<AgentMessageContent>,
 514    pub tool_results: IndexMap<LanguageModelToolUseId, LanguageModelToolResult>,
 515}
 516
 517#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 518pub enum AgentMessageContent {
 519    Text(String),
 520    Thinking {
 521        text: String,
 522        signature: Option<String>,
 523    },
 524    RedactedThinking(String),
 525    ToolUse(LanguageModelToolUse),
 526}
 527
 528pub trait TerminalHandle {
 529    fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId>;
 530    fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse>;
 531    fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>>;
 532}
 533
 534pub trait ThreadEnvironment {
 535    fn create_terminal(
 536        &self,
 537        command: String,
 538        cwd: Option<PathBuf>,
 539        output_byte_limit: Option<u64>,
 540        cx: &mut AsyncApp,
 541    ) -> Task<Result<Rc<dyn TerminalHandle>>>;
 542}
 543
 544#[derive(Debug)]
 545pub enum ThreadEvent {
 546    UserMessage(UserMessage),
 547    AgentText(String),
 548    AgentThinking(String),
 549    ToolCall(acp::ToolCall),
 550    ToolCallUpdate(acp_thread::ToolCallUpdate),
 551    ToolCallAuthorization(ToolCallAuthorization),
 552    Retry(acp_thread::RetryStatus),
 553    Stop(acp::StopReason),
 554}
 555
 556#[derive(Debug)]
 557pub struct NewTerminal {
 558    pub command: String,
 559    pub output_byte_limit: Option<u64>,
 560    pub cwd: Option<PathBuf>,
 561    pub response: oneshot::Sender<Result<Entity<acp_thread::Terminal>>>,
 562}
 563
 564#[derive(Debug)]
 565pub struct ToolCallAuthorization {
 566    pub tool_call: acp::ToolCallUpdate,
 567    pub options: Vec<acp::PermissionOption>,
 568    pub response: oneshot::Sender<acp::PermissionOptionId>,
 569}
 570
 571#[derive(Debug, thiserror::Error)]
 572enum CompletionError {
 573    #[error("max tokens")]
 574    MaxTokens,
 575    #[error("refusal")]
 576    Refusal,
 577    #[error(transparent)]
 578    Other(#[from] anyhow::Error),
 579}
 580
 581pub struct Thread {
 582    id: acp::SessionId,
 583    prompt_id: PromptId,
 584    updated_at: DateTime<Utc>,
 585    title: Option<SharedString>,
 586    pending_title_generation: Option<Task<()>>,
 587    summary: Option<SharedString>,
 588    messages: Vec<Message>,
 589    user_store: Entity<UserStore>,
 590    completion_mode: CompletionMode,
 591    /// Holds the task that handles agent interaction until the end of the turn.
 592    /// Survives across multiple requests as the model performs tool calls and
 593    /// we run tools, report their results.
 594    running_turn: Option<RunningTurn>,
 595    pending_message: Option<AgentMessage>,
 596    tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
 597    tool_use_limit_reached: bool,
 598    request_token_usage: HashMap<UserMessageId, language_model::TokenUsage>,
 599    #[allow(unused)]
 600    cumulative_token_usage: TokenUsage,
 601    #[allow(unused)]
 602    initial_project_snapshot: Shared<Task<Option<Arc<ProjectSnapshot>>>>,
 603    context_server_registry: Entity<ContextServerRegistry>,
 604    profile_id: AgentProfileId,
 605    project_context: Entity<ProjectContext>,
 606    templates: Arc<Templates>,
 607    model: Option<Arc<dyn LanguageModel>>,
 608    summarization_model: Option<Arc<dyn LanguageModel>>,
 609    prompt_capabilities_tx: watch::Sender<acp::PromptCapabilities>,
 610    pub(crate) prompt_capabilities_rx: watch::Receiver<acp::PromptCapabilities>,
 611    pub(crate) project: Entity<Project>,
 612    pub(crate) action_log: Entity<ActionLog>,
 613}
 614
 615impl Thread {
 616    fn prompt_capabilities(model: Option<&dyn LanguageModel>) -> acp::PromptCapabilities {
 617        let image = model.map_or(true, |model| model.supports_images());
 618        acp::PromptCapabilities {
 619            meta: None,
 620            image,
 621            audio: false,
 622            embedded_context: true,
 623        }
 624    }
 625
 626    pub fn new(
 627        project: Entity<Project>,
 628        project_context: Entity<ProjectContext>,
 629        context_server_registry: Entity<ContextServerRegistry>,
 630        templates: Arc<Templates>,
 631        model: Option<Arc<dyn LanguageModel>>,
 632        cx: &mut Context<Self>,
 633    ) -> Self {
 634        let profile_id = AgentSettings::get_global(cx).default_profile.clone();
 635        let action_log = cx.new(|_cx| ActionLog::new(project.clone()));
 636        let (prompt_capabilities_tx, prompt_capabilities_rx) =
 637            watch::channel(Self::prompt_capabilities(model.as_deref()));
 638        Self {
 639            id: acp::SessionId(uuid::Uuid::new_v4().to_string().into()),
 640            prompt_id: PromptId::new(),
 641            updated_at: Utc::now(),
 642            title: None,
 643            pending_title_generation: None,
 644            summary: None,
 645            messages: Vec::new(),
 646            user_store: project.read(cx).user_store(),
 647            completion_mode: AgentSettings::get_global(cx).preferred_completion_mode,
 648            running_turn: None,
 649            pending_message: None,
 650            tools: BTreeMap::default(),
 651            tool_use_limit_reached: false,
 652            request_token_usage: HashMap::default(),
 653            cumulative_token_usage: TokenUsage::default(),
 654            initial_project_snapshot: {
 655                let project_snapshot = Self::project_snapshot(project.clone(), cx);
 656                cx.foreground_executor()
 657                    .spawn(async move { Some(project_snapshot.await) })
 658                    .shared()
 659            },
 660            context_server_registry,
 661            profile_id,
 662            project_context,
 663            templates,
 664            model,
 665            summarization_model: None,
 666            prompt_capabilities_tx,
 667            prompt_capabilities_rx,
 668            project,
 669            action_log,
 670        }
 671    }
 672
 673    pub fn id(&self) -> &acp::SessionId {
 674        &self.id
 675    }
 676
 677    pub fn replay(
 678        &mut self,
 679        cx: &mut Context<Self>,
 680    ) -> mpsc::UnboundedReceiver<Result<ThreadEvent>> {
 681        let (tx, rx) = mpsc::unbounded();
 682        let stream = ThreadEventStream(tx);
 683        for message in &self.messages {
 684            match message {
 685                Message::User(user_message) => stream.send_user_message(user_message),
 686                Message::Agent(assistant_message) => {
 687                    for content in &assistant_message.content {
 688                        match content {
 689                            AgentMessageContent::Text(text) => stream.send_text(text),
 690                            AgentMessageContent::Thinking { text, .. } => {
 691                                stream.send_thinking(text)
 692                            }
 693                            AgentMessageContent::RedactedThinking(_) => {}
 694                            AgentMessageContent::ToolUse(tool_use) => {
 695                                self.replay_tool_call(
 696                                    tool_use,
 697                                    assistant_message.tool_results.get(&tool_use.id),
 698                                    &stream,
 699                                    cx,
 700                                );
 701                            }
 702                        }
 703                    }
 704                }
 705                Message::Resume => {}
 706            }
 707        }
 708        rx
 709    }
 710
 711    fn replay_tool_call(
 712        &self,
 713        tool_use: &LanguageModelToolUse,
 714        tool_result: Option<&LanguageModelToolResult>,
 715        stream: &ThreadEventStream,
 716        cx: &mut Context<Self>,
 717    ) {
 718        let tool = self.tools.get(tool_use.name.as_ref()).cloned().or_else(|| {
 719            self.context_server_registry
 720                .read(cx)
 721                .servers()
 722                .find_map(|(_, tools)| {
 723                    if let Some(tool) = tools.get(tool_use.name.as_ref()) {
 724                        Some(tool.clone())
 725                    } else {
 726                        None
 727                    }
 728                })
 729        });
 730
 731        let Some(tool) = tool else {
 732            stream
 733                .0
 734                .unbounded_send(Ok(ThreadEvent::ToolCall(acp::ToolCall {
 735                    meta: None,
 736                    id: acp::ToolCallId(tool_use.id.to_string().into()),
 737                    title: tool_use.name.to_string(),
 738                    kind: acp::ToolKind::Other,
 739                    status: acp::ToolCallStatus::Failed,
 740                    content: Vec::new(),
 741                    locations: Vec::new(),
 742                    raw_input: Some(tool_use.input.clone()),
 743                    raw_output: None,
 744                })))
 745                .ok();
 746            return;
 747        };
 748
 749        let title = tool.initial_title(tool_use.input.clone(), cx);
 750        let kind = tool.kind();
 751        stream.send_tool_call(&tool_use.id, title, kind, tool_use.input.clone());
 752
 753        let output = tool_result
 754            .as_ref()
 755            .and_then(|result| result.output.clone());
 756        if let Some(output) = output.clone() {
 757            let tool_event_stream = ToolCallEventStream::new(
 758                tool_use.id.clone(),
 759                stream.clone(),
 760                Some(self.project.read(cx).fs().clone()),
 761            );
 762            tool.replay(tool_use.input.clone(), output, tool_event_stream, cx)
 763                .log_err();
 764        }
 765
 766        stream.update_tool_call_fields(
 767            &tool_use.id,
 768            acp::ToolCallUpdateFields {
 769                status: Some(
 770                    tool_result
 771                        .as_ref()
 772                        .map_or(acp::ToolCallStatus::Failed, |result| {
 773                            if result.is_error {
 774                                acp::ToolCallStatus::Failed
 775                            } else {
 776                                acp::ToolCallStatus::Completed
 777                            }
 778                        }),
 779                ),
 780                raw_output: output,
 781                ..Default::default()
 782            },
 783        );
 784    }
 785
 786    pub fn from_db(
 787        id: acp::SessionId,
 788        db_thread: DbThread,
 789        project: Entity<Project>,
 790        project_context: Entity<ProjectContext>,
 791        context_server_registry: Entity<ContextServerRegistry>,
 792        action_log: Entity<ActionLog>,
 793        templates: Arc<Templates>,
 794        cx: &mut Context<Self>,
 795    ) -> Self {
 796        let profile_id = db_thread
 797            .profile
 798            .unwrap_or_else(|| AgentSettings::get_global(cx).default_profile.clone());
 799        let model = LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
 800            db_thread
 801                .model
 802                .and_then(|model| {
 803                    let model = SelectedModel {
 804                        provider: model.provider.clone().into(),
 805                        model: model.model.into(),
 806                    };
 807                    registry.select_model(&model, cx)
 808                })
 809                .or_else(|| registry.default_model())
 810                .map(|model| model.model)
 811        });
 812        let (prompt_capabilities_tx, prompt_capabilities_rx) =
 813            watch::channel(Self::prompt_capabilities(model.as_deref()));
 814
 815        Self {
 816            id,
 817            prompt_id: PromptId::new(),
 818            title: if db_thread.title.is_empty() {
 819                None
 820            } else {
 821                Some(db_thread.title.clone())
 822            },
 823            pending_title_generation: None,
 824            summary: db_thread.detailed_summary,
 825            messages: db_thread.messages,
 826            user_store: project.read(cx).user_store(),
 827            completion_mode: db_thread.completion_mode.unwrap_or_default(),
 828            running_turn: None,
 829            pending_message: None,
 830            tools: BTreeMap::default(),
 831            tool_use_limit_reached: false,
 832            request_token_usage: db_thread.request_token_usage.clone(),
 833            cumulative_token_usage: db_thread.cumulative_token_usage,
 834            initial_project_snapshot: Task::ready(db_thread.initial_project_snapshot).shared(),
 835            context_server_registry,
 836            profile_id,
 837            project_context,
 838            templates,
 839            model,
 840            summarization_model: None,
 841            project,
 842            action_log,
 843            updated_at: db_thread.updated_at,
 844            prompt_capabilities_tx,
 845            prompt_capabilities_rx,
 846        }
 847    }
 848
 849    pub fn to_db(&self, cx: &App) -> Task<DbThread> {
 850        let initial_project_snapshot = self.initial_project_snapshot.clone();
 851        let mut thread = DbThread {
 852            title: self.title(),
 853            messages: self.messages.clone(),
 854            updated_at: self.updated_at,
 855            detailed_summary: self.summary.clone(),
 856            initial_project_snapshot: None,
 857            cumulative_token_usage: self.cumulative_token_usage,
 858            request_token_usage: self.request_token_usage.clone(),
 859            model: self.model.as_ref().map(|model| DbLanguageModel {
 860                provider: model.provider_id().to_string(),
 861                model: model.name().0.to_string(),
 862            }),
 863            completion_mode: Some(self.completion_mode),
 864            profile: Some(self.profile_id.clone()),
 865        };
 866
 867        cx.background_spawn(async move {
 868            let initial_project_snapshot = initial_project_snapshot.await;
 869            thread.initial_project_snapshot = initial_project_snapshot;
 870            thread
 871        })
 872    }
 873
 874    /// Create a snapshot of the current project state including git information and unsaved buffers.
 875    fn project_snapshot(
 876        project: Entity<Project>,
 877        cx: &mut Context<Self>,
 878    ) -> Task<Arc<agent::thread::ProjectSnapshot>> {
 879        let git_store = project.read(cx).git_store().clone();
 880        let worktree_snapshots: Vec<_> = project
 881            .read(cx)
 882            .visible_worktrees(cx)
 883            .map(|worktree| Self::worktree_snapshot(worktree, git_store.clone(), cx))
 884            .collect();
 885
 886        cx.spawn(async move |_, _| {
 887            let worktree_snapshots = futures::future::join_all(worktree_snapshots).await;
 888
 889            Arc::new(ProjectSnapshot {
 890                worktree_snapshots,
 891                timestamp: Utc::now(),
 892            })
 893        })
 894    }
 895
 896    fn worktree_snapshot(
 897        worktree: Entity<project::Worktree>,
 898        git_store: Entity<GitStore>,
 899        cx: &App,
 900    ) -> Task<agent::thread::WorktreeSnapshot> {
 901        cx.spawn(async move |cx| {
 902            // Get worktree path and snapshot
 903            let worktree_info = cx.update(|app_cx| {
 904                let worktree = worktree.read(app_cx);
 905                let path = worktree.abs_path().to_string_lossy().into_owned();
 906                let snapshot = worktree.snapshot();
 907                (path, snapshot)
 908            });
 909
 910            let Ok((worktree_path, _snapshot)) = worktree_info else {
 911                return WorktreeSnapshot {
 912                    worktree_path: String::new(),
 913                    git_state: None,
 914                };
 915            };
 916
 917            let git_state = git_store
 918                .update(cx, |git_store, cx| {
 919                    git_store
 920                        .repositories()
 921                        .values()
 922                        .find(|repo| {
 923                            repo.read(cx)
 924                                .abs_path_to_repo_path(&worktree.read(cx).abs_path())
 925                                .is_some()
 926                        })
 927                        .cloned()
 928                })
 929                .ok()
 930                .flatten()
 931                .map(|repo| {
 932                    repo.update(cx, |repo, _| {
 933                        let current_branch =
 934                            repo.branch.as_ref().map(|branch| branch.name().to_owned());
 935                        repo.send_job(None, |state, _| async move {
 936                            let RepositoryState::Local { backend, .. } = state else {
 937                                return GitState {
 938                                    remote_url: None,
 939                                    head_sha: None,
 940                                    current_branch,
 941                                    diff: None,
 942                                };
 943                            };
 944
 945                            let remote_url = backend.remote_url("origin");
 946                            let head_sha = backend.head_sha().await;
 947                            let diff = backend.diff(DiffType::HeadToWorktree).await.ok();
 948
 949                            GitState {
 950                                remote_url,
 951                                head_sha,
 952                                current_branch,
 953                                diff,
 954                            }
 955                        })
 956                    })
 957                });
 958
 959            let git_state = match git_state {
 960                Some(git_state) => match git_state.ok() {
 961                    Some(git_state) => git_state.await.ok(),
 962                    None => None,
 963                },
 964                None => None,
 965            };
 966
 967            WorktreeSnapshot {
 968                worktree_path,
 969                git_state,
 970            }
 971        })
 972    }
 973
 974    pub fn project_context(&self) -> &Entity<ProjectContext> {
 975        &self.project_context
 976    }
 977
 978    pub fn project(&self) -> &Entity<Project> {
 979        &self.project
 980    }
 981
 982    pub fn action_log(&self) -> &Entity<ActionLog> {
 983        &self.action_log
 984    }
 985
 986    pub fn is_empty(&self) -> bool {
 987        self.messages.is_empty() && self.title.is_none()
 988    }
 989
 990    pub fn model(&self) -> Option<&Arc<dyn LanguageModel>> {
 991        self.model.as_ref()
 992    }
 993
 994    pub fn set_model(&mut self, model: Arc<dyn LanguageModel>, cx: &mut Context<Self>) {
 995        let old_usage = self.latest_token_usage();
 996        self.model = Some(model);
 997        let new_caps = Self::prompt_capabilities(self.model.as_deref());
 998        let new_usage = self.latest_token_usage();
 999        if old_usage != new_usage {
1000            cx.emit(TokenUsageUpdated(new_usage));
1001        }
1002        self.prompt_capabilities_tx.send(new_caps).log_err();
1003        cx.notify()
1004    }
1005
1006    pub fn summarization_model(&self) -> Option<&Arc<dyn LanguageModel>> {
1007        self.summarization_model.as_ref()
1008    }
1009
1010    pub fn set_summarization_model(
1011        &mut self,
1012        model: Option<Arc<dyn LanguageModel>>,
1013        cx: &mut Context<Self>,
1014    ) {
1015        self.summarization_model = model;
1016        cx.notify()
1017    }
1018
1019    pub fn completion_mode(&self) -> CompletionMode {
1020        self.completion_mode
1021    }
1022
1023    pub fn set_completion_mode(&mut self, mode: CompletionMode, cx: &mut Context<Self>) {
1024        let old_usage = self.latest_token_usage();
1025        self.completion_mode = mode;
1026        let new_usage = self.latest_token_usage();
1027        if old_usage != new_usage {
1028            cx.emit(TokenUsageUpdated(new_usage));
1029        }
1030        cx.notify()
1031    }
1032
1033    #[cfg(any(test, feature = "test-support"))]
1034    pub fn last_message(&self) -> Option<Message> {
1035        if let Some(message) = self.pending_message.clone() {
1036            Some(Message::Agent(message))
1037        } else {
1038            self.messages.last().cloned()
1039        }
1040    }
1041
1042    pub fn add_default_tools(
1043        &mut self,
1044        environment: Rc<dyn ThreadEnvironment>,
1045        cx: &mut Context<Self>,
1046    ) {
1047        let language_registry = self.project.read(cx).languages().clone();
1048        self.add_tool(CopyPathTool::new(self.project.clone()));
1049        self.add_tool(CreateDirectoryTool::new(self.project.clone()));
1050        self.add_tool(DeletePathTool::new(
1051            self.project.clone(),
1052            self.action_log.clone(),
1053        ));
1054        self.add_tool(DiagnosticsTool::new(self.project.clone()));
1055        self.add_tool(EditFileTool::new(
1056            self.project.clone(),
1057            cx.weak_entity(),
1058            language_registry,
1059        ));
1060        self.add_tool(FetchTool::new(self.project.read(cx).client().http_client()));
1061        self.add_tool(FindPathTool::new(self.project.clone()));
1062        self.add_tool(GrepTool::new(self.project.clone()));
1063        self.add_tool(ListDirectoryTool::new(self.project.clone()));
1064        self.add_tool(MovePathTool::new(self.project.clone()));
1065        self.add_tool(NowTool);
1066        self.add_tool(OpenTool::new(self.project.clone()));
1067        self.add_tool(ReadFileTool::new(
1068            self.project.clone(),
1069            self.action_log.clone(),
1070        ));
1071        self.add_tool(TerminalTool::new(self.project.clone(), environment));
1072        self.add_tool(ThinkingTool);
1073        self.add_tool(WebSearchTool);
1074    }
1075
1076    pub fn add_tool<T: AgentTool>(&mut self, tool: T) {
1077        self.tools.insert(T::name().into(), tool.erase());
1078    }
1079
1080    pub fn remove_tool(&mut self, name: &str) -> bool {
1081        self.tools.remove(name).is_some()
1082    }
1083
1084    pub fn profile(&self) -> &AgentProfileId {
1085        &self.profile_id
1086    }
1087
1088    pub fn set_profile(&mut self, profile_id: AgentProfileId) {
1089        self.profile_id = profile_id;
1090    }
1091
1092    pub fn cancel(&mut self, cx: &mut Context<Self>) {
1093        if let Some(running_turn) = self.running_turn.take() {
1094            running_turn.cancel();
1095        }
1096        self.flush_pending_message(cx);
1097    }
1098
1099    fn update_token_usage(&mut self, update: language_model::TokenUsage, cx: &mut Context<Self>) {
1100        let Some(last_user_message) = self.last_user_message() else {
1101            return;
1102        };
1103
1104        self.request_token_usage
1105            .insert(last_user_message.id.clone(), update);
1106        cx.emit(TokenUsageUpdated(self.latest_token_usage()));
1107        cx.notify();
1108    }
1109
1110    pub fn truncate(&mut self, message_id: UserMessageId, cx: &mut Context<Self>) -> Result<()> {
1111        self.cancel(cx);
1112        let Some(position) = self.messages.iter().position(
1113            |msg| matches!(msg, Message::User(UserMessage { id, .. }) if id == &message_id),
1114        ) else {
1115            return Err(anyhow!("Message not found"));
1116        };
1117
1118        for message in self.messages.drain(position..) {
1119            match message {
1120                Message::User(message) => {
1121                    self.request_token_usage.remove(&message.id);
1122                }
1123                Message::Agent(_) | Message::Resume => {}
1124            }
1125        }
1126        self.summary = None;
1127        cx.notify();
1128        Ok(())
1129    }
1130
1131    pub fn latest_token_usage(&self) -> Option<acp_thread::TokenUsage> {
1132        let last_user_message = self.last_user_message()?;
1133        let tokens = self.request_token_usage.get(&last_user_message.id)?;
1134        let model = self.model.clone()?;
1135
1136        Some(acp_thread::TokenUsage {
1137            max_tokens: model.max_token_count_for_mode(self.completion_mode.into()),
1138            used_tokens: tokens.total_tokens(),
1139        })
1140    }
1141
1142    pub fn resume(
1143        &mut self,
1144        cx: &mut Context<Self>,
1145    ) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>> {
1146        self.messages.push(Message::Resume);
1147        cx.notify();
1148
1149        log::debug!("Total messages in thread: {}", self.messages.len());
1150        self.run_turn(cx)
1151    }
1152
1153    /// Sending a message results in the model streaming a response, which could include tool calls.
1154    /// After calling tools, the model will stops and waits for any outstanding tool calls to be completed and their results sent.
1155    /// The returned channel will report all the occurrences in which the model stops before erroring or ending its turn.
1156    pub fn send<T>(
1157        &mut self,
1158        id: UserMessageId,
1159        content: impl IntoIterator<Item = T>,
1160        cx: &mut Context<Self>,
1161    ) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>
1162    where
1163        T: Into<UserMessageContent>,
1164    {
1165        let model = self.model().context("No language model configured")?;
1166
1167        log::info!("Thread::send called with model: {}", model.name().0);
1168        self.advance_prompt_id();
1169
1170        let content = content.into_iter().map(Into::into).collect::<Vec<_>>();
1171        log::debug!("Thread::send content: {:?}", content);
1172
1173        self.messages
1174            .push(Message::User(UserMessage { id, content }));
1175        cx.notify();
1176
1177        log::debug!("Total messages in thread: {}", self.messages.len());
1178        self.run_turn(cx)
1179    }
1180
1181    fn run_turn(
1182        &mut self,
1183        cx: &mut Context<Self>,
1184    ) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>> {
1185        self.cancel(cx);
1186
1187        let model = self.model.clone().context("No language model configured")?;
1188        let profile = AgentSettings::get_global(cx)
1189            .profiles
1190            .get(&self.profile_id)
1191            .context("Profile not found")?;
1192        let (events_tx, events_rx) = mpsc::unbounded::<Result<ThreadEvent>>();
1193        let event_stream = ThreadEventStream(events_tx);
1194        let message_ix = self.messages.len().saturating_sub(1);
1195        self.tool_use_limit_reached = false;
1196        self.summary = None;
1197        self.running_turn = Some(RunningTurn {
1198            event_stream: event_stream.clone(),
1199            tools: self.enabled_tools(profile, &model, cx),
1200            _task: cx.spawn(async move |this, cx| {
1201                log::debug!("Starting agent turn execution");
1202
1203                let turn_result = Self::run_turn_internal(&this, model, &event_stream, cx).await;
1204                _ = this.update(cx, |this, cx| this.flush_pending_message(cx));
1205
1206                match turn_result {
1207                    Ok(()) => {
1208                        log::debug!("Turn execution completed");
1209                        event_stream.send_stop(acp::StopReason::EndTurn);
1210                    }
1211                    Err(error) => {
1212                        log::error!("Turn execution failed: {:?}", error);
1213                        match error.downcast::<CompletionError>() {
1214                            Ok(CompletionError::Refusal) => {
1215                                event_stream.send_stop(acp::StopReason::Refusal);
1216                                _ = this.update(cx, |this, _| this.messages.truncate(message_ix));
1217                            }
1218                            Ok(CompletionError::MaxTokens) => {
1219                                event_stream.send_stop(acp::StopReason::MaxTokens);
1220                            }
1221                            Ok(CompletionError::Other(error)) | Err(error) => {
1222                                event_stream.send_error(error);
1223                            }
1224                        }
1225                    }
1226                }
1227
1228                _ = this.update(cx, |this, _| this.running_turn.take());
1229            }),
1230        });
1231        Ok(events_rx)
1232    }
1233
1234    async fn run_turn_internal(
1235        this: &WeakEntity<Self>,
1236        model: Arc<dyn LanguageModel>,
1237        event_stream: &ThreadEventStream,
1238        cx: &mut AsyncApp,
1239    ) -> Result<()> {
1240        let mut attempt = 0;
1241        let mut intent = CompletionIntent::UserPrompt;
1242        loop {
1243            let request =
1244                this.update(cx, |this, cx| this.build_completion_request(intent, cx))??;
1245
1246            telemetry::event!(
1247                "Agent Thread Completion",
1248                thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
1249                prompt_id = this.read_with(cx, |this, _| this.prompt_id.to_string())?,
1250                model = model.telemetry_id(),
1251                model_provider = model.provider_id().to_string(),
1252                attempt
1253            );
1254
1255            log::debug!("Calling model.stream_completion, attempt {}", attempt);
1256
1257            let (mut events, mut error) = match model.stream_completion(request, cx).await {
1258                Ok(events) => (events, None),
1259                Err(err) => (stream::empty().boxed(), Some(err)),
1260            };
1261            let mut tool_results = FuturesUnordered::new();
1262            while let Some(event) = events.next().await {
1263                log::trace!("Received completion event: {:?}", event);
1264                match event {
1265                    Ok(event) => {
1266                        tool_results.extend(this.update(cx, |this, cx| {
1267                            this.handle_completion_event(event, event_stream, cx)
1268                        })??);
1269                    }
1270                    Err(err) => {
1271                        error = Some(err);
1272                        break;
1273                    }
1274                }
1275            }
1276
1277            let end_turn = tool_results.is_empty();
1278            while let Some(tool_result) = tool_results.next().await {
1279                log::debug!("Tool finished {:?}", tool_result);
1280
1281                event_stream.update_tool_call_fields(
1282                    &tool_result.tool_use_id,
1283                    acp::ToolCallUpdateFields {
1284                        status: Some(if tool_result.is_error {
1285                            acp::ToolCallStatus::Failed
1286                        } else {
1287                            acp::ToolCallStatus::Completed
1288                        }),
1289                        raw_output: tool_result.output.clone(),
1290                        ..Default::default()
1291                    },
1292                );
1293                this.update(cx, |this, _cx| {
1294                    this.pending_message()
1295                        .tool_results
1296                        .insert(tool_result.tool_use_id.clone(), tool_result);
1297                })?;
1298            }
1299
1300            this.update(cx, |this, cx| {
1301                this.flush_pending_message(cx);
1302                if this.title.is_none() && this.pending_title_generation.is_none() {
1303                    this.generate_title(cx);
1304                }
1305            })?;
1306
1307            if let Some(error) = error {
1308                attempt += 1;
1309                let retry = this.update(cx, |this, cx| {
1310                    let user_store = this.user_store.read(cx);
1311                    this.handle_completion_error(error, attempt, user_store.plan())
1312                })??;
1313                let timer = cx.background_executor().timer(retry.duration);
1314                event_stream.send_retry(retry);
1315                timer.await;
1316                this.update(cx, |this, _cx| {
1317                    if let Some(Message::Agent(message)) = this.messages.last() {
1318                        if message.tool_results.is_empty() {
1319                            intent = CompletionIntent::UserPrompt;
1320                            this.messages.push(Message::Resume);
1321                        }
1322                    }
1323                })?;
1324            } else if this.read_with(cx, |this, _| this.tool_use_limit_reached)? {
1325                return Err(language_model::ToolUseLimitReachedError.into());
1326            } else if end_turn {
1327                return Ok(());
1328            } else {
1329                intent = CompletionIntent::ToolResults;
1330                attempt = 0;
1331            }
1332        }
1333    }
1334
1335    fn handle_completion_error(
1336        &mut self,
1337        error: LanguageModelCompletionError,
1338        attempt: u8,
1339        plan: Option<Plan>,
1340    ) -> Result<acp_thread::RetryStatus> {
1341        let Some(model) = self.model.as_ref() else {
1342            return Err(anyhow!(error));
1343        };
1344
1345        let auto_retry = if model.provider_id() == ZED_CLOUD_PROVIDER_ID {
1346            match plan {
1347                Some(Plan::V2(_)) => true,
1348                Some(Plan::V1(_)) => self.completion_mode == CompletionMode::Burn,
1349                None => false,
1350            }
1351        } else {
1352            true
1353        };
1354
1355        if !auto_retry {
1356            return Err(anyhow!(error));
1357        }
1358
1359        let Some(strategy) = Self::retry_strategy_for(&error) else {
1360            return Err(anyhow!(error));
1361        };
1362
1363        let max_attempts = match &strategy {
1364            RetryStrategy::ExponentialBackoff { max_attempts, .. } => *max_attempts,
1365            RetryStrategy::Fixed { max_attempts, .. } => *max_attempts,
1366        };
1367
1368        if attempt > max_attempts {
1369            return Err(anyhow!(error));
1370        }
1371
1372        let delay = match &strategy {
1373            RetryStrategy::ExponentialBackoff { initial_delay, .. } => {
1374                let delay_secs = initial_delay.as_secs() * 2u64.pow((attempt - 1) as u32);
1375                Duration::from_secs(delay_secs)
1376            }
1377            RetryStrategy::Fixed { delay, .. } => *delay,
1378        };
1379        log::debug!("Retry attempt {attempt} with delay {delay:?}");
1380
1381        Ok(acp_thread::RetryStatus {
1382            last_error: error.to_string().into(),
1383            attempt: attempt as usize,
1384            max_attempts: max_attempts as usize,
1385            started_at: Instant::now(),
1386            duration: delay,
1387        })
1388    }
1389
1390    /// A helper method that's called on every streamed completion event.
1391    /// Returns an optional tool result task, which the main agentic loop will
1392    /// send back to the model when it resolves.
1393    fn handle_completion_event(
1394        &mut self,
1395        event: LanguageModelCompletionEvent,
1396        event_stream: &ThreadEventStream,
1397        cx: &mut Context<Self>,
1398    ) -> Result<Option<Task<LanguageModelToolResult>>> {
1399        log::trace!("Handling streamed completion event: {:?}", event);
1400        use LanguageModelCompletionEvent::*;
1401
1402        match event {
1403            StartMessage { .. } => {
1404                self.flush_pending_message(cx);
1405                self.pending_message = Some(AgentMessage::default());
1406            }
1407            Text(new_text) => self.handle_text_event(new_text, event_stream, cx),
1408            Thinking { text, signature } => {
1409                self.handle_thinking_event(text, signature, event_stream, cx)
1410            }
1411            RedactedThinking { data } => self.handle_redacted_thinking_event(data, cx),
1412            ToolUse(tool_use) => {
1413                return Ok(self.handle_tool_use_event(tool_use, event_stream, cx));
1414            }
1415            ToolUseJsonParseError {
1416                id,
1417                tool_name,
1418                raw_input,
1419                json_parse_error,
1420            } => {
1421                return Ok(Some(Task::ready(
1422                    self.handle_tool_use_json_parse_error_event(
1423                        id,
1424                        tool_name,
1425                        raw_input,
1426                        json_parse_error,
1427                    ),
1428                )));
1429            }
1430            UsageUpdate(usage) => {
1431                telemetry::event!(
1432                    "Agent Thread Completion Usage Updated",
1433                    thread_id = self.id.to_string(),
1434                    prompt_id = self.prompt_id.to_string(),
1435                    model = self.model.as_ref().map(|m| m.telemetry_id()),
1436                    model_provider = self.model.as_ref().map(|m| m.provider_id().to_string()),
1437                    input_tokens = usage.input_tokens,
1438                    output_tokens = usage.output_tokens,
1439                    cache_creation_input_tokens = usage.cache_creation_input_tokens,
1440                    cache_read_input_tokens = usage.cache_read_input_tokens,
1441                );
1442                self.update_token_usage(usage, cx);
1443            }
1444            StatusUpdate(CompletionRequestStatus::UsageUpdated { amount, limit }) => {
1445                self.update_model_request_usage(amount, limit, cx);
1446            }
1447            StatusUpdate(
1448                CompletionRequestStatus::Started
1449                | CompletionRequestStatus::Queued { .. }
1450                | CompletionRequestStatus::Failed { .. },
1451            ) => {}
1452            StatusUpdate(CompletionRequestStatus::ToolUseLimitReached) => {
1453                self.tool_use_limit_reached = true;
1454            }
1455            Stop(StopReason::Refusal) => return Err(CompletionError::Refusal.into()),
1456            Stop(StopReason::MaxTokens) => return Err(CompletionError::MaxTokens.into()),
1457            Stop(StopReason::ToolUse | StopReason::EndTurn) => {}
1458        }
1459
1460        Ok(None)
1461    }
1462
1463    fn handle_text_event(
1464        &mut self,
1465        new_text: String,
1466        event_stream: &ThreadEventStream,
1467        cx: &mut Context<Self>,
1468    ) {
1469        event_stream.send_text(&new_text);
1470
1471        let last_message = self.pending_message();
1472        if let Some(AgentMessageContent::Text(text)) = last_message.content.last_mut() {
1473            text.push_str(&new_text);
1474        } else {
1475            last_message
1476                .content
1477                .push(AgentMessageContent::Text(new_text));
1478        }
1479
1480        cx.notify();
1481    }
1482
1483    fn handle_thinking_event(
1484        &mut self,
1485        new_text: String,
1486        new_signature: Option<String>,
1487        event_stream: &ThreadEventStream,
1488        cx: &mut Context<Self>,
1489    ) {
1490        event_stream.send_thinking(&new_text);
1491
1492        let last_message = self.pending_message();
1493        if let Some(AgentMessageContent::Thinking { text, signature }) =
1494            last_message.content.last_mut()
1495        {
1496            text.push_str(&new_text);
1497            *signature = new_signature.or(signature.take());
1498        } else {
1499            last_message.content.push(AgentMessageContent::Thinking {
1500                text: new_text,
1501                signature: new_signature,
1502            });
1503        }
1504
1505        cx.notify();
1506    }
1507
1508    fn handle_redacted_thinking_event(&mut self, data: String, cx: &mut Context<Self>) {
1509        let last_message = self.pending_message();
1510        last_message
1511            .content
1512            .push(AgentMessageContent::RedactedThinking(data));
1513        cx.notify();
1514    }
1515
1516    fn handle_tool_use_event(
1517        &mut self,
1518        tool_use: LanguageModelToolUse,
1519        event_stream: &ThreadEventStream,
1520        cx: &mut Context<Self>,
1521    ) -> Option<Task<LanguageModelToolResult>> {
1522        cx.notify();
1523
1524        let tool = self.tool(tool_use.name.as_ref());
1525        let mut title = SharedString::from(&tool_use.name);
1526        let mut kind = acp::ToolKind::Other;
1527        if let Some(tool) = tool.as_ref() {
1528            title = tool.initial_title(tool_use.input.clone(), cx);
1529            kind = tool.kind();
1530        }
1531
1532        // Ensure the last message ends in the current tool use
1533        let last_message = self.pending_message();
1534        let push_new_tool_use = last_message.content.last_mut().is_none_or(|content| {
1535            if let AgentMessageContent::ToolUse(last_tool_use) = content {
1536                if last_tool_use.id == tool_use.id {
1537                    *last_tool_use = tool_use.clone();
1538                    false
1539                } else {
1540                    true
1541                }
1542            } else {
1543                true
1544            }
1545        });
1546
1547        if push_new_tool_use {
1548            event_stream.send_tool_call(&tool_use.id, title, kind, tool_use.input.clone());
1549            last_message
1550                .content
1551                .push(AgentMessageContent::ToolUse(tool_use.clone()));
1552        } else {
1553            event_stream.update_tool_call_fields(
1554                &tool_use.id,
1555                acp::ToolCallUpdateFields {
1556                    title: Some(title.into()),
1557                    kind: Some(kind),
1558                    raw_input: Some(tool_use.input.clone()),
1559                    ..Default::default()
1560                },
1561            );
1562        }
1563
1564        if !tool_use.is_input_complete {
1565            return None;
1566        }
1567
1568        let Some(tool) = tool else {
1569            let content = format!("No tool named {} exists", tool_use.name);
1570            return Some(Task::ready(LanguageModelToolResult {
1571                content: LanguageModelToolResultContent::Text(Arc::from(content)),
1572                tool_use_id: tool_use.id,
1573                tool_name: tool_use.name,
1574                is_error: true,
1575                output: None,
1576            }));
1577        };
1578
1579        let fs = self.project.read(cx).fs().clone();
1580        let tool_event_stream =
1581            ToolCallEventStream::new(tool_use.id.clone(), event_stream.clone(), Some(fs));
1582        tool_event_stream.update_fields(acp::ToolCallUpdateFields {
1583            status: Some(acp::ToolCallStatus::InProgress),
1584            ..Default::default()
1585        });
1586        let supports_images = self.model().is_some_and(|model| model.supports_images());
1587        let tool_result = tool.run(tool_use.input, tool_event_stream, cx);
1588        log::debug!("Running tool {}", tool_use.name);
1589        Some(cx.foreground_executor().spawn(async move {
1590            let tool_result = tool_result.await.and_then(|output| {
1591                if let LanguageModelToolResultContent::Image(_) = &output.llm_output
1592                    && !supports_images
1593                {
1594                    return Err(anyhow!(
1595                        "Attempted to read an image, but this model doesn't support it.",
1596                    ));
1597                }
1598                Ok(output)
1599            });
1600
1601            match tool_result {
1602                Ok(output) => LanguageModelToolResult {
1603                    tool_use_id: tool_use.id,
1604                    tool_name: tool_use.name,
1605                    is_error: false,
1606                    content: output.llm_output,
1607                    output: Some(output.raw_output),
1608                },
1609                Err(error) => LanguageModelToolResult {
1610                    tool_use_id: tool_use.id,
1611                    tool_name: tool_use.name,
1612                    is_error: true,
1613                    content: LanguageModelToolResultContent::Text(Arc::from(error.to_string())),
1614                    output: Some(error.to_string().into()),
1615                },
1616            }
1617        }))
1618    }
1619
1620    fn handle_tool_use_json_parse_error_event(
1621        &mut self,
1622        tool_use_id: LanguageModelToolUseId,
1623        tool_name: Arc<str>,
1624        raw_input: Arc<str>,
1625        json_parse_error: String,
1626    ) -> LanguageModelToolResult {
1627        let tool_output = format!("Error parsing input JSON: {json_parse_error}");
1628        LanguageModelToolResult {
1629            tool_use_id,
1630            tool_name,
1631            is_error: true,
1632            content: LanguageModelToolResultContent::Text(tool_output.into()),
1633            output: Some(serde_json::Value::String(raw_input.to_string())),
1634        }
1635    }
1636
1637    fn update_model_request_usage(&self, amount: usize, limit: UsageLimit, cx: &mut Context<Self>) {
1638        self.project
1639            .read(cx)
1640            .user_store()
1641            .update(cx, |user_store, cx| {
1642                user_store.update_model_request_usage(
1643                    ModelRequestUsage(RequestUsage {
1644                        amount: amount as i32,
1645                        limit,
1646                    }),
1647                    cx,
1648                )
1649            });
1650    }
1651
1652    pub fn title(&self) -> SharedString {
1653        self.title.clone().unwrap_or("New Thread".into())
1654    }
1655
1656    pub fn summary(&mut self, cx: &mut Context<Self>) -> Task<Result<SharedString>> {
1657        if let Some(summary) = self.summary.as_ref() {
1658            return Task::ready(Ok(summary.clone()));
1659        }
1660        let Some(model) = self.summarization_model.clone() else {
1661            return Task::ready(Err(anyhow!("No summarization model available")));
1662        };
1663        let mut request = LanguageModelRequest {
1664            intent: Some(CompletionIntent::ThreadContextSummarization),
1665            temperature: AgentSettings::temperature_for_model(&model, cx),
1666            ..Default::default()
1667        };
1668
1669        for message in &self.messages {
1670            request.messages.extend(message.to_request());
1671        }
1672
1673        request.messages.push(LanguageModelRequestMessage {
1674            role: Role::User,
1675            content: vec![SUMMARIZE_THREAD_DETAILED_PROMPT.into()],
1676            cache: false,
1677        });
1678        cx.spawn(async move |this, cx| {
1679            let mut summary = String::new();
1680            let mut messages = model.stream_completion(request, cx).await?;
1681            while let Some(event) = messages.next().await {
1682                let event = event?;
1683                let text = match event {
1684                    LanguageModelCompletionEvent::Text(text) => text,
1685                    LanguageModelCompletionEvent::StatusUpdate(
1686                        CompletionRequestStatus::UsageUpdated { amount, limit },
1687                    ) => {
1688                        this.update(cx, |thread, cx| {
1689                            thread.update_model_request_usage(amount, limit, cx);
1690                        })?;
1691                        continue;
1692                    }
1693                    _ => continue,
1694                };
1695
1696                let mut lines = text.lines();
1697                summary.extend(lines.next());
1698            }
1699
1700            log::debug!("Setting summary: {}", summary);
1701            let summary = SharedString::from(summary);
1702
1703            this.update(cx, |this, cx| {
1704                this.summary = Some(summary.clone());
1705                cx.notify()
1706            })?;
1707
1708            Ok(summary)
1709        })
1710    }
1711
1712    fn generate_title(&mut self, cx: &mut Context<Self>) {
1713        let Some(model) = self.summarization_model.clone() else {
1714            return;
1715        };
1716
1717        log::debug!(
1718            "Generating title with model: {:?}",
1719            self.summarization_model.as_ref().map(|model| model.name())
1720        );
1721        let mut request = LanguageModelRequest {
1722            intent: Some(CompletionIntent::ThreadSummarization),
1723            temperature: AgentSettings::temperature_for_model(&model, cx),
1724            ..Default::default()
1725        };
1726
1727        for message in &self.messages {
1728            request.messages.extend(message.to_request());
1729        }
1730
1731        request.messages.push(LanguageModelRequestMessage {
1732            role: Role::User,
1733            content: vec![SUMMARIZE_THREAD_PROMPT.into()],
1734            cache: false,
1735        });
1736        self.pending_title_generation = Some(cx.spawn(async move |this, cx| {
1737            let mut title = String::new();
1738
1739            let generate = async {
1740                let mut messages = model.stream_completion(request, cx).await?;
1741                while let Some(event) = messages.next().await {
1742                    let event = event?;
1743                    let text = match event {
1744                        LanguageModelCompletionEvent::Text(text) => text,
1745                        LanguageModelCompletionEvent::StatusUpdate(
1746                            CompletionRequestStatus::UsageUpdated { amount, limit },
1747                        ) => {
1748                            this.update(cx, |thread, cx| {
1749                                thread.update_model_request_usage(amount, limit, cx);
1750                            })?;
1751                            continue;
1752                        }
1753                        _ => continue,
1754                    };
1755
1756                    let mut lines = text.lines();
1757                    title.extend(lines.next());
1758
1759                    // Stop if the LLM generated multiple lines.
1760                    if lines.next().is_some() {
1761                        break;
1762                    }
1763                }
1764                anyhow::Ok(())
1765            };
1766
1767            if generate.await.context("failed to generate title").is_ok() {
1768                _ = this.update(cx, |this, cx| this.set_title(title.into(), cx));
1769            }
1770            _ = this.update(cx, |this, _| this.pending_title_generation = None);
1771        }));
1772    }
1773
1774    pub fn set_title(&mut self, title: SharedString, cx: &mut Context<Self>) {
1775        self.pending_title_generation = None;
1776        if Some(&title) != self.title.as_ref() {
1777            self.title = Some(title);
1778            cx.emit(TitleUpdated);
1779            cx.notify();
1780        }
1781    }
1782
1783    fn last_user_message(&self) -> Option<&UserMessage> {
1784        self.messages
1785            .iter()
1786            .rev()
1787            .find_map(|message| match message {
1788                Message::User(user_message) => Some(user_message),
1789                Message::Agent(_) => None,
1790                Message::Resume => None,
1791            })
1792    }
1793
1794    fn pending_message(&mut self) -> &mut AgentMessage {
1795        self.pending_message.get_or_insert_default()
1796    }
1797
1798    fn flush_pending_message(&mut self, cx: &mut Context<Self>) {
1799        let Some(mut message) = self.pending_message.take() else {
1800            return;
1801        };
1802
1803        if message.content.is_empty() {
1804            return;
1805        }
1806
1807        for content in &message.content {
1808            let AgentMessageContent::ToolUse(tool_use) = content else {
1809                continue;
1810            };
1811
1812            if !message.tool_results.contains_key(&tool_use.id) {
1813                message.tool_results.insert(
1814                    tool_use.id.clone(),
1815                    LanguageModelToolResult {
1816                        tool_use_id: tool_use.id.clone(),
1817                        tool_name: tool_use.name.clone(),
1818                        is_error: true,
1819                        content: LanguageModelToolResultContent::Text(TOOL_CANCELED_MESSAGE.into()),
1820                        output: None,
1821                    },
1822                );
1823            }
1824        }
1825
1826        self.messages.push(Message::Agent(message));
1827        self.updated_at = Utc::now();
1828        self.summary = None;
1829        cx.notify()
1830    }
1831
1832    pub(crate) fn build_completion_request(
1833        &self,
1834        completion_intent: CompletionIntent,
1835        cx: &App,
1836    ) -> Result<LanguageModelRequest> {
1837        let model = self.model().context("No language model configured")?;
1838        let tools = if let Some(turn) = self.running_turn.as_ref() {
1839            turn.tools
1840                .iter()
1841                .filter_map(|(tool_name, tool)| {
1842                    log::trace!("Including tool: {}", tool_name);
1843                    Some(LanguageModelRequestTool {
1844                        name: tool_name.to_string(),
1845                        description: tool.description().to_string(),
1846                        input_schema: tool.input_schema(model.tool_input_format()).log_err()?,
1847                    })
1848                })
1849                .collect::<Vec<_>>()
1850        } else {
1851            Vec::new()
1852        };
1853
1854        log::debug!("Building completion request");
1855        log::debug!("Completion intent: {:?}", completion_intent);
1856        log::debug!("Completion mode: {:?}", self.completion_mode);
1857
1858        let messages = self.build_request_messages(cx);
1859        log::debug!("Request will include {} messages", messages.len());
1860        log::debug!("Request includes {} tools", tools.len());
1861
1862        let request = LanguageModelRequest {
1863            thread_id: Some(self.id.to_string()),
1864            prompt_id: Some(self.prompt_id.to_string()),
1865            intent: Some(completion_intent),
1866            mode: Some(self.completion_mode.into()),
1867            messages,
1868            tools,
1869            tool_choice: None,
1870            stop: Vec::new(),
1871            temperature: AgentSettings::temperature_for_model(model, cx),
1872            thinking_allowed: true,
1873        };
1874
1875        log::debug!("Completion request built successfully");
1876        Ok(request)
1877    }
1878
1879    fn enabled_tools(
1880        &self,
1881        profile: &AgentProfileSettings,
1882        model: &Arc<dyn LanguageModel>,
1883        cx: &App,
1884    ) -> BTreeMap<SharedString, Arc<dyn AnyAgentTool>> {
1885        fn truncate(tool_name: &SharedString) -> SharedString {
1886            if tool_name.len() > MAX_TOOL_NAME_LENGTH {
1887                let mut truncated = tool_name.to_string();
1888                truncated.truncate(MAX_TOOL_NAME_LENGTH);
1889                truncated.into()
1890            } else {
1891                tool_name.clone()
1892            }
1893        }
1894
1895        let mut tools = self
1896            .tools
1897            .iter()
1898            .filter_map(|(tool_name, tool)| {
1899                if tool.supported_provider(&model.provider_id())
1900                    && profile.is_tool_enabled(tool_name)
1901                {
1902                    Some((truncate(tool_name), tool.clone()))
1903                } else {
1904                    None
1905                }
1906            })
1907            .collect::<BTreeMap<_, _>>();
1908
1909        let mut context_server_tools = Vec::new();
1910        let mut seen_tools = tools.keys().cloned().collect::<HashSet<_>>();
1911        let mut duplicate_tool_names = HashSet::default();
1912        for (server_id, server_tools) in self.context_server_registry.read(cx).servers() {
1913            for (tool_name, tool) in server_tools {
1914                if profile.is_context_server_tool_enabled(&server_id.0, &tool_name) {
1915                    let tool_name = truncate(tool_name);
1916                    if !seen_tools.insert(tool_name.clone()) {
1917                        duplicate_tool_names.insert(tool_name.clone());
1918                    }
1919                    context_server_tools.push((server_id.clone(), tool_name, tool.clone()));
1920                }
1921            }
1922        }
1923
1924        // When there are duplicate tool names, disambiguate by prefixing them
1925        // with the server ID. In the rare case there isn't enough space for the
1926        // disambiguated tool name, keep only the last tool with this name.
1927        for (server_id, tool_name, tool) in context_server_tools {
1928            if duplicate_tool_names.contains(&tool_name) {
1929                let available = MAX_TOOL_NAME_LENGTH.saturating_sub(tool_name.len());
1930                if available >= 2 {
1931                    let mut disambiguated = server_id.0.to_string();
1932                    disambiguated.truncate(available - 1);
1933                    disambiguated.push('_');
1934                    disambiguated.push_str(&tool_name);
1935                    tools.insert(disambiguated.into(), tool.clone());
1936                } else {
1937                    tools.insert(tool_name, tool.clone());
1938                }
1939            } else {
1940                tools.insert(tool_name, tool.clone());
1941            }
1942        }
1943
1944        tools
1945    }
1946
1947    fn tool(&self, name: &str) -> Option<Arc<dyn AnyAgentTool>> {
1948        self.running_turn.as_ref()?.tools.get(name).cloned()
1949    }
1950
1951    fn build_request_messages(&self, cx: &App) -> Vec<LanguageModelRequestMessage> {
1952        log::trace!(
1953            "Building request messages from {} thread messages",
1954            self.messages.len()
1955        );
1956
1957        let system_prompt = SystemPromptTemplate {
1958            project: self.project_context.read(cx),
1959            available_tools: self.tools.keys().cloned().collect(),
1960        }
1961        .render(&self.templates)
1962        .context("failed to build system prompt")
1963        .expect("Invalid template");
1964        let mut messages = vec![LanguageModelRequestMessage {
1965            role: Role::System,
1966            content: vec![system_prompt.into()],
1967            cache: false,
1968        }];
1969        for message in &self.messages {
1970            messages.extend(message.to_request());
1971        }
1972
1973        if let Some(last_message) = messages.last_mut() {
1974            last_message.cache = true;
1975        }
1976
1977        if let Some(message) = self.pending_message.as_ref() {
1978            messages.extend(message.to_request());
1979        }
1980
1981        messages
1982    }
1983
1984    pub fn to_markdown(&self) -> String {
1985        let mut markdown = String::new();
1986        for (ix, message) in self.messages.iter().enumerate() {
1987            if ix > 0 {
1988                markdown.push('\n');
1989            }
1990            markdown.push_str(&message.to_markdown());
1991        }
1992
1993        if let Some(message) = self.pending_message.as_ref() {
1994            markdown.push('\n');
1995            markdown.push_str(&message.to_markdown());
1996        }
1997
1998        markdown
1999    }
2000
2001    fn advance_prompt_id(&mut self) {
2002        self.prompt_id = PromptId::new();
2003    }
2004
2005    fn retry_strategy_for(error: &LanguageModelCompletionError) -> Option<RetryStrategy> {
2006        use LanguageModelCompletionError::*;
2007        use http_client::StatusCode;
2008
2009        // General strategy here:
2010        // - If retrying won't help (e.g. invalid API key or payload too large), return None so we don't retry at all.
2011        // - If it's a time-based issue (e.g. server overloaded, rate limit exceeded), retry up to 4 times with exponential backoff.
2012        // - If it's an issue that *might* be fixed by retrying (e.g. internal server error), retry up to 3 times.
2013        match error {
2014            HttpResponseError {
2015                status_code: StatusCode::TOO_MANY_REQUESTS,
2016                ..
2017            } => Some(RetryStrategy::ExponentialBackoff {
2018                initial_delay: BASE_RETRY_DELAY,
2019                max_attempts: MAX_RETRY_ATTEMPTS,
2020            }),
2021            ServerOverloaded { retry_after, .. } | RateLimitExceeded { retry_after, .. } => {
2022                Some(RetryStrategy::Fixed {
2023                    delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2024                    max_attempts: MAX_RETRY_ATTEMPTS,
2025                })
2026            }
2027            UpstreamProviderError {
2028                status,
2029                retry_after,
2030                ..
2031            } => match *status {
2032                StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => {
2033                    Some(RetryStrategy::Fixed {
2034                        delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2035                        max_attempts: MAX_RETRY_ATTEMPTS,
2036                    })
2037                }
2038                StatusCode::INTERNAL_SERVER_ERROR => Some(RetryStrategy::Fixed {
2039                    delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2040                    // Internal Server Error could be anything, retry up to 3 times.
2041                    max_attempts: 3,
2042                }),
2043                status => {
2044                    // There is no StatusCode variant for the unofficial HTTP 529 ("The service is overloaded"),
2045                    // but we frequently get them in practice. See https://http.dev/529
2046                    if status.as_u16() == 529 {
2047                        Some(RetryStrategy::Fixed {
2048                            delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2049                            max_attempts: MAX_RETRY_ATTEMPTS,
2050                        })
2051                    } else {
2052                        Some(RetryStrategy::Fixed {
2053                            delay: retry_after.unwrap_or(BASE_RETRY_DELAY),
2054                            max_attempts: 2,
2055                        })
2056                    }
2057                }
2058            },
2059            ApiInternalServerError { .. } => Some(RetryStrategy::Fixed {
2060                delay: BASE_RETRY_DELAY,
2061                max_attempts: 3,
2062            }),
2063            ApiReadResponseError { .. }
2064            | HttpSend { .. }
2065            | DeserializeResponse { .. }
2066            | BadRequestFormat { .. } => Some(RetryStrategy::Fixed {
2067                delay: BASE_RETRY_DELAY,
2068                max_attempts: 3,
2069            }),
2070            // Retrying these errors definitely shouldn't help.
2071            HttpResponseError {
2072                status_code:
2073                    StatusCode::PAYLOAD_TOO_LARGE | StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED,
2074                ..
2075            }
2076            | AuthenticationError { .. }
2077            | PermissionError { .. }
2078            | NoApiKey { .. }
2079            | ApiEndpointNotFound { .. }
2080            | PromptTooLarge { .. } => None,
2081            // These errors might be transient, so retry them
2082            SerializeRequest { .. } | BuildRequestBody { .. } => Some(RetryStrategy::Fixed {
2083                delay: BASE_RETRY_DELAY,
2084                max_attempts: 1,
2085            }),
2086            // Retry all other 4xx and 5xx errors once.
2087            HttpResponseError { status_code, .. }
2088                if status_code.is_client_error() || status_code.is_server_error() =>
2089            {
2090                Some(RetryStrategy::Fixed {
2091                    delay: BASE_RETRY_DELAY,
2092                    max_attempts: 3,
2093                })
2094            }
2095            Other(err)
2096                if err.is::<language_model::PaymentRequiredError>()
2097                    || err.is::<language_model::ModelRequestLimitReachedError>() =>
2098            {
2099                // Retrying won't help for Payment Required or Model Request Limit errors (where
2100                // the user must upgrade to usage-based billing to get more requests, or else wait
2101                // for a significant amount of time for the request limit to reset).
2102                None
2103            }
2104            // Conservatively assume that any other errors are non-retryable
2105            HttpResponseError { .. } | Other(..) => Some(RetryStrategy::Fixed {
2106                delay: BASE_RETRY_DELAY,
2107                max_attempts: 2,
2108            }),
2109        }
2110    }
2111}
2112
2113struct RunningTurn {
2114    /// Holds the task that handles agent interaction until the end of the turn.
2115    /// Survives across multiple requests as the model performs tool calls and
2116    /// we run tools, report their results.
2117    _task: Task<()>,
2118    /// The current event stream for the running turn. Used to report a final
2119    /// cancellation event if we cancel the turn.
2120    event_stream: ThreadEventStream,
2121    /// The tools that were enabled for this turn.
2122    tools: BTreeMap<SharedString, Arc<dyn AnyAgentTool>>,
2123}
2124
2125impl RunningTurn {
2126    fn cancel(self) {
2127        log::debug!("Cancelling in progress turn");
2128        self.event_stream.send_canceled();
2129    }
2130}
2131
2132pub struct TokenUsageUpdated(pub Option<acp_thread::TokenUsage>);
2133
2134impl EventEmitter<TokenUsageUpdated> for Thread {}
2135
2136pub struct TitleUpdated;
2137
2138impl EventEmitter<TitleUpdated> for Thread {}
2139
2140pub trait AgentTool
2141where
2142    Self: 'static + Sized,
2143{
2144    type Input: for<'de> Deserialize<'de> + Serialize + JsonSchema;
2145    type Output: for<'de> Deserialize<'de> + Serialize + Into<LanguageModelToolResultContent>;
2146
2147    fn name() -> &'static str;
2148
2149    fn description(&self) -> SharedString {
2150        let schema = schemars::schema_for!(Self::Input);
2151        SharedString::new(
2152            schema
2153                .get("description")
2154                .and_then(|description| description.as_str())
2155                .unwrap_or_default(),
2156        )
2157    }
2158
2159    fn kind() -> acp::ToolKind;
2160
2161    /// The initial tool title to display. Can be updated during the tool run.
2162    fn initial_title(
2163        &self,
2164        input: Result<Self::Input, serde_json::Value>,
2165        cx: &mut App,
2166    ) -> SharedString;
2167
2168    /// Returns the JSON schema that describes the tool's input.
2169    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Schema {
2170        crate::tool_schema::root_schema_for::<Self::Input>(format)
2171    }
2172
2173    /// Some tools rely on a provider for the underlying billing or other reasons.
2174    /// Allow the tool to check if they are compatible, or should be filtered out.
2175    fn supported_provider(&self, _provider: &LanguageModelProviderId) -> bool {
2176        true
2177    }
2178
2179    /// Runs the tool with the provided input.
2180    fn run(
2181        self: Arc<Self>,
2182        input: Self::Input,
2183        event_stream: ToolCallEventStream,
2184        cx: &mut App,
2185    ) -> Task<Result<Self::Output>>;
2186
2187    /// Emits events for a previous execution of the tool.
2188    fn replay(
2189        &self,
2190        _input: Self::Input,
2191        _output: Self::Output,
2192        _event_stream: ToolCallEventStream,
2193        _cx: &mut App,
2194    ) -> Result<()> {
2195        Ok(())
2196    }
2197
2198    fn erase(self) -> Arc<dyn AnyAgentTool> {
2199        Arc::new(Erased(Arc::new(self)))
2200    }
2201}
2202
2203pub struct Erased<T>(T);
2204
2205pub struct AgentToolOutput {
2206    pub llm_output: LanguageModelToolResultContent,
2207    pub raw_output: serde_json::Value,
2208}
2209
2210pub trait AnyAgentTool {
2211    fn name(&self) -> SharedString;
2212    fn description(&self) -> SharedString;
2213    fn kind(&self) -> acp::ToolKind;
2214    fn initial_title(&self, input: serde_json::Value, _cx: &mut App) -> SharedString;
2215    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value>;
2216    fn supported_provider(&self, _provider: &LanguageModelProviderId) -> bool {
2217        true
2218    }
2219    fn run(
2220        self: Arc<Self>,
2221        input: serde_json::Value,
2222        event_stream: ToolCallEventStream,
2223        cx: &mut App,
2224    ) -> Task<Result<AgentToolOutput>>;
2225    fn replay(
2226        &self,
2227        input: serde_json::Value,
2228        output: serde_json::Value,
2229        event_stream: ToolCallEventStream,
2230        cx: &mut App,
2231    ) -> Result<()>;
2232}
2233
2234impl<T> AnyAgentTool for Erased<Arc<T>>
2235where
2236    T: AgentTool,
2237{
2238    fn name(&self) -> SharedString {
2239        T::name().into()
2240    }
2241
2242    fn description(&self) -> SharedString {
2243        self.0.description()
2244    }
2245
2246    fn kind(&self) -> agent_client_protocol::ToolKind {
2247        T::kind()
2248    }
2249
2250    fn initial_title(&self, input: serde_json::Value, _cx: &mut App) -> SharedString {
2251        let parsed_input = serde_json::from_value(input.clone()).map_err(|_| input);
2252        self.0.initial_title(parsed_input, _cx)
2253    }
2254
2255    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
2256        let mut json = serde_json::to_value(self.0.input_schema(format))?;
2257        adapt_schema_to_format(&mut json, format)?;
2258        Ok(json)
2259    }
2260
2261    fn supported_provider(&self, provider: &LanguageModelProviderId) -> bool {
2262        self.0.supported_provider(provider)
2263    }
2264
2265    fn run(
2266        self: Arc<Self>,
2267        input: serde_json::Value,
2268        event_stream: ToolCallEventStream,
2269        cx: &mut App,
2270    ) -> Task<Result<AgentToolOutput>> {
2271        cx.spawn(async move |cx| {
2272            let input = serde_json::from_value(input)?;
2273            let output = cx
2274                .update(|cx| self.0.clone().run(input, event_stream, cx))?
2275                .await?;
2276            let raw_output = serde_json::to_value(&output)?;
2277            Ok(AgentToolOutput {
2278                llm_output: output.into(),
2279                raw_output,
2280            })
2281        })
2282    }
2283
2284    fn replay(
2285        &self,
2286        input: serde_json::Value,
2287        output: serde_json::Value,
2288        event_stream: ToolCallEventStream,
2289        cx: &mut App,
2290    ) -> Result<()> {
2291        let input = serde_json::from_value(input)?;
2292        let output = serde_json::from_value(output)?;
2293        self.0.replay(input, output, event_stream, cx)
2294    }
2295}
2296
2297#[derive(Clone)]
2298struct ThreadEventStream(mpsc::UnboundedSender<Result<ThreadEvent>>);
2299
2300impl ThreadEventStream {
2301    fn send_user_message(&self, message: &UserMessage) {
2302        self.0
2303            .unbounded_send(Ok(ThreadEvent::UserMessage(message.clone())))
2304            .ok();
2305    }
2306
2307    fn send_text(&self, text: &str) {
2308        self.0
2309            .unbounded_send(Ok(ThreadEvent::AgentText(text.to_string())))
2310            .ok();
2311    }
2312
2313    fn send_thinking(&self, text: &str) {
2314        self.0
2315            .unbounded_send(Ok(ThreadEvent::AgentThinking(text.to_string())))
2316            .ok();
2317    }
2318
2319    fn send_tool_call(
2320        &self,
2321        id: &LanguageModelToolUseId,
2322        title: SharedString,
2323        kind: acp::ToolKind,
2324        input: serde_json::Value,
2325    ) {
2326        self.0
2327            .unbounded_send(Ok(ThreadEvent::ToolCall(Self::initial_tool_call(
2328                id,
2329                title.to_string(),
2330                kind,
2331                input,
2332            ))))
2333            .ok();
2334    }
2335
2336    fn initial_tool_call(
2337        id: &LanguageModelToolUseId,
2338        title: String,
2339        kind: acp::ToolKind,
2340        input: serde_json::Value,
2341    ) -> acp::ToolCall {
2342        acp::ToolCall {
2343            meta: None,
2344            id: acp::ToolCallId(id.to_string().into()),
2345            title,
2346            kind,
2347            status: acp::ToolCallStatus::Pending,
2348            content: vec![],
2349            locations: vec![],
2350            raw_input: Some(input),
2351            raw_output: None,
2352        }
2353    }
2354
2355    fn update_tool_call_fields(
2356        &self,
2357        tool_use_id: &LanguageModelToolUseId,
2358        fields: acp::ToolCallUpdateFields,
2359    ) {
2360        self.0
2361            .unbounded_send(Ok(ThreadEvent::ToolCallUpdate(
2362                acp::ToolCallUpdate {
2363                    meta: None,
2364                    id: acp::ToolCallId(tool_use_id.to_string().into()),
2365                    fields,
2366                }
2367                .into(),
2368            )))
2369            .ok();
2370    }
2371
2372    fn send_retry(&self, status: acp_thread::RetryStatus) {
2373        self.0.unbounded_send(Ok(ThreadEvent::Retry(status))).ok();
2374    }
2375
2376    fn send_stop(&self, reason: acp::StopReason) {
2377        self.0.unbounded_send(Ok(ThreadEvent::Stop(reason))).ok();
2378    }
2379
2380    fn send_canceled(&self) {
2381        self.0
2382            .unbounded_send(Ok(ThreadEvent::Stop(acp::StopReason::Cancelled)))
2383            .ok();
2384    }
2385
2386    fn send_error(&self, error: impl Into<anyhow::Error>) {
2387        self.0.unbounded_send(Err(error.into())).ok();
2388    }
2389}
2390
2391#[derive(Clone)]
2392pub struct ToolCallEventStream {
2393    tool_use_id: LanguageModelToolUseId,
2394    stream: ThreadEventStream,
2395    fs: Option<Arc<dyn Fs>>,
2396}
2397
2398impl ToolCallEventStream {
2399    #[cfg(test)]
2400    pub fn test() -> (Self, ToolCallEventStreamReceiver) {
2401        let (events_tx, events_rx) = mpsc::unbounded::<Result<ThreadEvent>>();
2402
2403        let stream = ToolCallEventStream::new("test_id".into(), ThreadEventStream(events_tx), None);
2404
2405        (stream, ToolCallEventStreamReceiver(events_rx))
2406    }
2407
2408    fn new(
2409        tool_use_id: LanguageModelToolUseId,
2410        stream: ThreadEventStream,
2411        fs: Option<Arc<dyn Fs>>,
2412    ) -> Self {
2413        Self {
2414            tool_use_id,
2415            stream,
2416            fs,
2417        }
2418    }
2419
2420    pub fn update_fields(&self, fields: acp::ToolCallUpdateFields) {
2421        self.stream
2422            .update_tool_call_fields(&self.tool_use_id, fields);
2423    }
2424
2425    pub fn update_diff(&self, diff: Entity<acp_thread::Diff>) {
2426        self.stream
2427            .0
2428            .unbounded_send(Ok(ThreadEvent::ToolCallUpdate(
2429                acp_thread::ToolCallUpdateDiff {
2430                    id: acp::ToolCallId(self.tool_use_id.to_string().into()),
2431                    diff,
2432                }
2433                .into(),
2434            )))
2435            .ok();
2436    }
2437
2438    pub fn authorize(&self, title: impl Into<String>, cx: &mut App) -> Task<Result<()>> {
2439        if agent_settings::AgentSettings::get_global(cx).always_allow_tool_actions {
2440            return Task::ready(Ok(()));
2441        }
2442
2443        let (response_tx, response_rx) = oneshot::channel();
2444        self.stream
2445            .0
2446            .unbounded_send(Ok(ThreadEvent::ToolCallAuthorization(
2447                ToolCallAuthorization {
2448                    tool_call: acp::ToolCallUpdate {
2449                        meta: None,
2450                        id: acp::ToolCallId(self.tool_use_id.to_string().into()),
2451                        fields: acp::ToolCallUpdateFields {
2452                            title: Some(title.into()),
2453                            ..Default::default()
2454                        },
2455                    },
2456                    options: vec![
2457                        acp::PermissionOption {
2458                            id: acp::PermissionOptionId("always_allow".into()),
2459                            name: "Always Allow".into(),
2460                            kind: acp::PermissionOptionKind::AllowAlways,
2461                            meta: None,
2462                        },
2463                        acp::PermissionOption {
2464                            id: acp::PermissionOptionId("allow".into()),
2465                            name: "Allow".into(),
2466                            kind: acp::PermissionOptionKind::AllowOnce,
2467                            meta: None,
2468                        },
2469                        acp::PermissionOption {
2470                            id: acp::PermissionOptionId("deny".into()),
2471                            name: "Deny".into(),
2472                            kind: acp::PermissionOptionKind::RejectOnce,
2473                            meta: None,
2474                        },
2475                    ],
2476                    response: response_tx,
2477                },
2478            )))
2479            .ok();
2480        let fs = self.fs.clone();
2481        cx.spawn(async move |cx| match response_rx.await?.0.as_ref() {
2482            "always_allow" => {
2483                if let Some(fs) = fs.clone() {
2484                    cx.update(|cx| {
2485                        update_settings_file(fs, cx, |settings, _| {
2486                            settings
2487                                .agent
2488                                .get_or_insert_default()
2489                                .set_always_allow_tool_actions(true);
2490                        });
2491                    })?;
2492                }
2493
2494                Ok(())
2495            }
2496            "allow" => Ok(()),
2497            _ => Err(anyhow!("Permission to run tool denied by user")),
2498        })
2499    }
2500}
2501
2502#[cfg(test)]
2503pub struct ToolCallEventStreamReceiver(mpsc::UnboundedReceiver<Result<ThreadEvent>>);
2504
2505#[cfg(test)]
2506impl ToolCallEventStreamReceiver {
2507    pub async fn expect_authorization(&mut self) -> ToolCallAuthorization {
2508        let event = self.0.next().await;
2509        if let Some(Ok(ThreadEvent::ToolCallAuthorization(auth))) = event {
2510            auth
2511        } else {
2512            panic!("Expected ToolCallAuthorization but got: {:?}", event);
2513        }
2514    }
2515
2516    pub async fn expect_update_fields(&mut self) -> acp::ToolCallUpdateFields {
2517        let event = self.0.next().await;
2518        if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(
2519            update,
2520        )))) = event
2521        {
2522            update.fields
2523        } else {
2524            panic!("Expected update fields but got: {:?}", event);
2525        }
2526    }
2527
2528    pub async fn expect_diff(&mut self) -> Entity<acp_thread::Diff> {
2529        let event = self.0.next().await;
2530        if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateDiff(
2531            update,
2532        )))) = event
2533        {
2534            update.diff
2535        } else {
2536            panic!("Expected diff but got: {:?}", event);
2537        }
2538    }
2539
2540    pub async fn expect_terminal(&mut self) -> Entity<acp_thread::Terminal> {
2541        let event = self.0.next().await;
2542        if let Some(Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateTerminal(
2543            update,
2544        )))) = event
2545        {
2546            update.terminal
2547        } else {
2548            panic!("Expected terminal but got: {:?}", event);
2549        }
2550    }
2551}
2552
2553#[cfg(test)]
2554impl std::ops::Deref for ToolCallEventStreamReceiver {
2555    type Target = mpsc::UnboundedReceiver<Result<ThreadEvent>>;
2556
2557    fn deref(&self) -> &Self::Target {
2558        &self.0
2559    }
2560}
2561
2562#[cfg(test)]
2563impl std::ops::DerefMut for ToolCallEventStreamReceiver {
2564    fn deref_mut(&mut self) -> &mut Self::Target {
2565        &mut self.0
2566    }
2567}
2568
2569impl From<&str> for UserMessageContent {
2570    fn from(text: &str) -> Self {
2571        Self::Text(text.into())
2572    }
2573}
2574
2575impl From<acp::ContentBlock> for UserMessageContent {
2576    fn from(value: acp::ContentBlock) -> Self {
2577        match value {
2578            acp::ContentBlock::Text(text_content) => Self::Text(text_content.text),
2579            acp::ContentBlock::Image(image_content) => Self::Image(convert_image(image_content)),
2580            acp::ContentBlock::Audio(_) => {
2581                // TODO
2582                Self::Text("[audio]".to_string())
2583            }
2584            acp::ContentBlock::ResourceLink(resource_link) => {
2585                match MentionUri::parse(&resource_link.uri) {
2586                    Ok(uri) => Self::Mention {
2587                        uri,
2588                        content: String::new(),
2589                    },
2590                    Err(err) => {
2591                        log::error!("Failed to parse mention link: {}", err);
2592                        Self::Text(format!("[{}]({})", resource_link.name, resource_link.uri))
2593                    }
2594                }
2595            }
2596            acp::ContentBlock::Resource(resource) => match resource.resource {
2597                acp::EmbeddedResourceResource::TextResourceContents(resource) => {
2598                    match MentionUri::parse(&resource.uri) {
2599                        Ok(uri) => Self::Mention {
2600                            uri,
2601                            content: resource.text,
2602                        },
2603                        Err(err) => {
2604                            log::error!("Failed to parse mention link: {}", err);
2605                            Self::Text(
2606                                MarkdownCodeBlock {
2607                                    tag: &resource.uri,
2608                                    text: &resource.text,
2609                                }
2610                                .to_string(),
2611                            )
2612                        }
2613                    }
2614                }
2615                acp::EmbeddedResourceResource::BlobResourceContents(_) => {
2616                    // TODO
2617                    Self::Text("[blob]".to_string())
2618                }
2619            },
2620        }
2621    }
2622}
2623
2624impl From<UserMessageContent> for acp::ContentBlock {
2625    fn from(content: UserMessageContent) -> Self {
2626        match content {
2627            UserMessageContent::Text(text) => acp::ContentBlock::Text(acp::TextContent {
2628                text,
2629                annotations: None,
2630                meta: None,
2631            }),
2632            UserMessageContent::Image(image) => acp::ContentBlock::Image(acp::ImageContent {
2633                data: image.source.to_string(),
2634                mime_type: "image/png".to_string(),
2635                meta: None,
2636                annotations: None,
2637                uri: None,
2638            }),
2639            UserMessageContent::Mention { uri, content } => {
2640                acp::ContentBlock::Resource(acp::EmbeddedResource {
2641                    meta: None,
2642                    resource: acp::EmbeddedResourceResource::TextResourceContents(
2643                        acp::TextResourceContents {
2644                            meta: None,
2645                            mime_type: None,
2646                            text: content,
2647                            uri: uri.to_uri().to_string(),
2648                        },
2649                    ),
2650                    annotations: None,
2651                })
2652            }
2653        }
2654    }
2655}
2656
2657fn convert_image(image_content: acp::ImageContent) -> LanguageModelImage {
2658    LanguageModelImage {
2659        source: image_content.data.into(),
2660        // TODO: make this optional?
2661        size: gpui::Size::new(0.into(), 0.into()),
2662    }
2663}