1mod connection;
   2mod diff;
   3mod mention;
   4mod terminal;
   5
   6use agent_settings::AgentSettings;
   7use collections::HashSet;
   8pub use connection::*;
   9pub use diff::*;
  10use language::language_settings::FormatOnSave;
  11pub use mention::*;
  12use project::lsp_store::{FormatTrigger, LspFormatTarget};
  13use serde::{Deserialize, Serialize};
  14use settings::Settings as _;
  15use task::{Shell, ShellBuilder};
  16pub use terminal::*;
  17
  18use action_log::ActionLog;
  19use agent_client_protocol::{self as acp};
  20use anyhow::{Context as _, Result, anyhow};
  21use editor::Bias;
  22use futures::{FutureExt, channel::oneshot, future::BoxFuture};
  23use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
  24use itertools::Itertools;
  25use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, ToPoint, text_diff};
  26use markdown::Markdown;
  27use project::{AgentLocation, Project, git_store::GitStoreCheckpoint};
  28use std::collections::HashMap;
  29use std::error::Error;
  30use std::fmt::{Formatter, Write};
  31use std::ops::Range;
  32use std::process::ExitStatus;
  33use std::rc::Rc;
  34use std::time::{Duration, Instant};
  35use std::{fmt::Display, mem, path::PathBuf, sync::Arc};
  36use ui::App;
  37use util::{ResultExt, get_default_system_shell};
  38use uuid::Uuid;
  39
  40#[derive(Debug)]
  41pub struct UserMessage {
  42    pub id: Option<UserMessageId>,
  43    pub content: ContentBlock,
  44    pub chunks: Vec<acp::ContentBlock>,
  45    pub checkpoint: Option<Checkpoint>,
  46}
  47
  48#[derive(Debug)]
  49pub struct Checkpoint {
  50    git_checkpoint: GitStoreCheckpoint,
  51    pub show: bool,
  52}
  53
  54impl UserMessage {
  55    fn to_markdown(&self, cx: &App) -> String {
  56        let mut markdown = String::new();
  57        if self
  58            .checkpoint
  59            .as_ref()
  60            .is_some_and(|checkpoint| checkpoint.show)
  61        {
  62            writeln!(markdown, "## User (checkpoint)").unwrap();
  63        } else {
  64            writeln!(markdown, "## User").unwrap();
  65        }
  66        writeln!(markdown).unwrap();
  67        writeln!(markdown, "{}", self.content.to_markdown(cx)).unwrap();
  68        writeln!(markdown).unwrap();
  69        markdown
  70    }
  71}
  72
  73#[derive(Debug, PartialEq)]
  74pub struct AssistantMessage {
  75    pub chunks: Vec<AssistantMessageChunk>,
  76}
  77
  78impl AssistantMessage {
  79    pub fn to_markdown(&self, cx: &App) -> String {
  80        format!(
  81            "## Assistant\n\n{}\n\n",
  82            self.chunks
  83                .iter()
  84                .map(|chunk| chunk.to_markdown(cx))
  85                .join("\n\n")
  86        )
  87    }
  88}
  89
  90#[derive(Debug, PartialEq)]
  91pub enum AssistantMessageChunk {
  92    Message { block: ContentBlock },
  93    Thought { block: ContentBlock },
  94}
  95
  96impl AssistantMessageChunk {
  97    pub fn from_str(chunk: &str, language_registry: &Arc<LanguageRegistry>, cx: &mut App) -> Self {
  98        Self::Message {
  99            block: ContentBlock::new(chunk.into(), language_registry, cx),
 100        }
 101    }
 102
 103    fn to_markdown(&self, cx: &App) -> String {
 104        match self {
 105            Self::Message { block } => block.to_markdown(cx).to_string(),
 106            Self::Thought { block } => {
 107                format!("<thinking>\n{}\n</thinking>", block.to_markdown(cx))
 108            }
 109        }
 110    }
 111}
 112
 113#[derive(Debug)]
 114pub enum AgentThreadEntry {
 115    UserMessage(UserMessage),
 116    AssistantMessage(AssistantMessage),
 117    ToolCall(ToolCall),
 118}
 119
 120impl AgentThreadEntry {
 121    pub fn to_markdown(&self, cx: &App) -> String {
 122        match self {
 123            Self::UserMessage(message) => message.to_markdown(cx),
 124            Self::AssistantMessage(message) => message.to_markdown(cx),
 125            Self::ToolCall(tool_call) => tool_call.to_markdown(cx),
 126        }
 127    }
 128
 129    pub fn user_message(&self) -> Option<&UserMessage> {
 130        if let AgentThreadEntry::UserMessage(message) = self {
 131            Some(message)
 132        } else {
 133            None
 134        }
 135    }
 136
 137    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 138        if let AgentThreadEntry::ToolCall(call) = self {
 139            itertools::Either::Left(call.diffs())
 140        } else {
 141            itertools::Either::Right(std::iter::empty())
 142        }
 143    }
 144
 145    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 146        if let AgentThreadEntry::ToolCall(call) = self {
 147            itertools::Either::Left(call.terminals())
 148        } else {
 149            itertools::Either::Right(std::iter::empty())
 150        }
 151    }
 152
 153    pub fn location(&self, ix: usize) -> Option<(acp::ToolCallLocation, AgentLocation)> {
 154        if let AgentThreadEntry::ToolCall(ToolCall {
 155            locations,
 156            resolved_locations,
 157            ..
 158        }) = self
 159        {
 160            Some((
 161                locations.get(ix)?.clone(),
 162                resolved_locations.get(ix)?.clone()?,
 163            ))
 164        } else {
 165            None
 166        }
 167    }
 168}
 169
 170#[derive(Debug)]
 171pub struct ToolCall {
 172    pub id: acp::ToolCallId,
 173    pub label: Entity<Markdown>,
 174    pub kind: acp::ToolKind,
 175    pub content: Vec<ToolCallContent>,
 176    pub status: ToolCallStatus,
 177    pub locations: Vec<acp::ToolCallLocation>,
 178    pub resolved_locations: Vec<Option<AgentLocation>>,
 179    pub raw_input: Option<serde_json::Value>,
 180    pub raw_output: Option<serde_json::Value>,
 181}
 182
 183impl ToolCall {
 184    fn from_acp(
 185        tool_call: acp::ToolCall,
 186        status: ToolCallStatus,
 187        language_registry: Arc<LanguageRegistry>,
 188        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 189        cx: &mut App,
 190    ) -> Result<Self> {
 191        let title = if let Some((first_line, _)) = tool_call.title.split_once("\n") {
 192            first_line.to_owned() + "…"
 193        } else {
 194            tool_call.title
 195        };
 196        let mut content = Vec::with_capacity(tool_call.content.len());
 197        for item in tool_call.content {
 198            content.push(ToolCallContent::from_acp(
 199                item,
 200                language_registry.clone(),
 201                terminals,
 202                cx,
 203            )?);
 204        }
 205
 206        let result = Self {
 207            id: tool_call.id,
 208            label: cx
 209                .new(|cx| Markdown::new(title.into(), Some(language_registry.clone()), None, cx)),
 210            kind: tool_call.kind,
 211            content,
 212            locations: tool_call.locations,
 213            resolved_locations: Vec::default(),
 214            status,
 215            raw_input: tool_call.raw_input,
 216            raw_output: tool_call.raw_output,
 217        };
 218        Ok(result)
 219    }
 220
 221    fn update_fields(
 222        &mut self,
 223        fields: acp::ToolCallUpdateFields,
 224        language_registry: Arc<LanguageRegistry>,
 225        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 226        cx: &mut App,
 227    ) -> Result<()> {
 228        let acp::ToolCallUpdateFields {
 229            kind,
 230            status,
 231            title,
 232            content,
 233            locations,
 234            raw_input,
 235            raw_output,
 236        } = fields;
 237
 238        if let Some(kind) = kind {
 239            self.kind = kind;
 240        }
 241
 242        if let Some(status) = status {
 243            self.status = status.into();
 244        }
 245
 246        if let Some(title) = title {
 247            self.label.update(cx, |label, cx| {
 248                if let Some((first_line, _)) = title.split_once("\n") {
 249                    label.replace(first_line.to_owned() + "…", cx)
 250                } else {
 251                    label.replace(title, cx);
 252                }
 253            });
 254        }
 255
 256        if let Some(content) = content {
 257            let new_content_len = content.len();
 258            let mut content = content.into_iter();
 259
 260            // Reuse existing content if we can
 261            for (old, new) in self.content.iter_mut().zip(content.by_ref()) {
 262                old.update_from_acp(new, language_registry.clone(), terminals, cx)?;
 263            }
 264            for new in content {
 265                self.content.push(ToolCallContent::from_acp(
 266                    new,
 267                    language_registry.clone(),
 268                    terminals,
 269                    cx,
 270                )?)
 271            }
 272            self.content.truncate(new_content_len);
 273        }
 274
 275        if let Some(locations) = locations {
 276            self.locations = locations;
 277        }
 278
 279        if let Some(raw_input) = raw_input {
 280            self.raw_input = Some(raw_input);
 281        }
 282
 283        if let Some(raw_output) = raw_output {
 284            if self.content.is_empty()
 285                && let Some(markdown) = markdown_for_raw_output(&raw_output, &language_registry, cx)
 286            {
 287                self.content
 288                    .push(ToolCallContent::ContentBlock(ContentBlock::Markdown {
 289                        markdown,
 290                    }));
 291            }
 292            self.raw_output = Some(raw_output);
 293        }
 294        Ok(())
 295    }
 296
 297    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 298        self.content.iter().filter_map(|content| match content {
 299            ToolCallContent::Diff(diff) => Some(diff),
 300            ToolCallContent::ContentBlock(_) => None,
 301            ToolCallContent::Terminal(_) => None,
 302        })
 303    }
 304
 305    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 306        self.content.iter().filter_map(|content| match content {
 307            ToolCallContent::Terminal(terminal) => Some(terminal),
 308            ToolCallContent::ContentBlock(_) => None,
 309            ToolCallContent::Diff(_) => None,
 310        })
 311    }
 312
 313    fn to_markdown(&self, cx: &App) -> String {
 314        let mut markdown = format!(
 315            "**Tool Call: {}**\nStatus: {}\n\n",
 316            self.label.read(cx).source(),
 317            self.status
 318        );
 319        for content in &self.content {
 320            markdown.push_str(content.to_markdown(cx).as_str());
 321            markdown.push_str("\n\n");
 322        }
 323        markdown
 324    }
 325
 326    async fn resolve_location(
 327        location: acp::ToolCallLocation,
 328        project: WeakEntity<Project>,
 329        cx: &mut AsyncApp,
 330    ) -> Option<AgentLocation> {
 331        let buffer = project
 332            .update(cx, |project, cx| {
 333                project
 334                    .project_path_for_absolute_path(&location.path, cx)
 335                    .map(|path| project.open_buffer(path, cx))
 336            })
 337            .ok()??;
 338        let buffer = buffer.await.log_err()?;
 339        let position = buffer
 340            .update(cx, |buffer, _| {
 341                if let Some(row) = location.line {
 342                    let snapshot = buffer.snapshot();
 343                    let column = snapshot.indent_size_for_line(row).len;
 344                    let point = snapshot.clip_point(Point::new(row, column), Bias::Left);
 345                    snapshot.anchor_before(point)
 346                } else {
 347                    Anchor::MIN
 348                }
 349            })
 350            .ok()?;
 351
 352        Some(AgentLocation {
 353            buffer: buffer.downgrade(),
 354            position,
 355        })
 356    }
 357
 358    fn resolve_locations(
 359        &self,
 360        project: Entity<Project>,
 361        cx: &mut App,
 362    ) -> Task<Vec<Option<AgentLocation>>> {
 363        let locations = self.locations.clone();
 364        project.update(cx, |_, cx| {
 365            cx.spawn(async move |project, cx| {
 366                let mut new_locations = Vec::new();
 367                for location in locations {
 368                    new_locations.push(Self::resolve_location(location, project.clone(), cx).await);
 369                }
 370                new_locations
 371            })
 372        })
 373    }
 374}
 375
 376#[derive(Debug)]
 377pub enum ToolCallStatus {
 378    /// The tool call hasn't started running yet, but we start showing it to
 379    /// the user.
 380    Pending,
 381    /// The tool call is waiting for confirmation from the user.
 382    WaitingForConfirmation {
 383        options: Vec<acp::PermissionOption>,
 384        respond_tx: oneshot::Sender<acp::PermissionOptionId>,
 385    },
 386    /// The tool call is currently running.
 387    InProgress,
 388    /// The tool call completed successfully.
 389    Completed,
 390    /// The tool call failed.
 391    Failed,
 392    /// The user rejected the tool call.
 393    Rejected,
 394    /// The user canceled generation so the tool call was canceled.
 395    Canceled,
 396}
 397
 398impl From<acp::ToolCallStatus> for ToolCallStatus {
 399    fn from(status: acp::ToolCallStatus) -> Self {
 400        match status {
 401            acp::ToolCallStatus::Pending => Self::Pending,
 402            acp::ToolCallStatus::InProgress => Self::InProgress,
 403            acp::ToolCallStatus::Completed => Self::Completed,
 404            acp::ToolCallStatus::Failed => Self::Failed,
 405        }
 406    }
 407}
 408
 409impl Display for ToolCallStatus {
 410    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 411        write!(
 412            f,
 413            "{}",
 414            match self {
 415                ToolCallStatus::Pending => "Pending",
 416                ToolCallStatus::WaitingForConfirmation { .. } => "Waiting for confirmation",
 417                ToolCallStatus::InProgress => "In Progress",
 418                ToolCallStatus::Completed => "Completed",
 419                ToolCallStatus::Failed => "Failed",
 420                ToolCallStatus::Rejected => "Rejected",
 421                ToolCallStatus::Canceled => "Canceled",
 422            }
 423        )
 424    }
 425}
 426
 427#[derive(Debug, PartialEq, Clone)]
 428pub enum ContentBlock {
 429    Empty,
 430    Markdown { markdown: Entity<Markdown> },
 431    ResourceLink { resource_link: acp::ResourceLink },
 432}
 433
 434impl ContentBlock {
 435    pub fn new(
 436        block: acp::ContentBlock,
 437        language_registry: &Arc<LanguageRegistry>,
 438        cx: &mut App,
 439    ) -> Self {
 440        let mut this = Self::Empty;
 441        this.append(block, language_registry, cx);
 442        this
 443    }
 444
 445    pub fn new_combined(
 446        blocks: impl IntoIterator<Item = acp::ContentBlock>,
 447        language_registry: Arc<LanguageRegistry>,
 448        cx: &mut App,
 449    ) -> Self {
 450        let mut this = Self::Empty;
 451        for block in blocks {
 452            this.append(block, &language_registry, cx);
 453        }
 454        this
 455    }
 456
 457    pub fn append(
 458        &mut self,
 459        block: acp::ContentBlock,
 460        language_registry: &Arc<LanguageRegistry>,
 461        cx: &mut App,
 462    ) {
 463        if matches!(self, ContentBlock::Empty)
 464            && let acp::ContentBlock::ResourceLink(resource_link) = block
 465        {
 466            *self = ContentBlock::ResourceLink { resource_link };
 467            return;
 468        }
 469
 470        let new_content = self.block_string_contents(block);
 471
 472        match self {
 473            ContentBlock::Empty => {
 474                *self = Self::create_markdown_block(new_content, language_registry, cx);
 475            }
 476            ContentBlock::Markdown { markdown } => {
 477                markdown.update(cx, |markdown, cx| markdown.append(&new_content, cx));
 478            }
 479            ContentBlock::ResourceLink { resource_link } => {
 480                let existing_content = Self::resource_link_md(&resource_link.uri);
 481                let combined = format!("{}\n{}", existing_content, new_content);
 482
 483                *self = Self::create_markdown_block(combined, language_registry, cx);
 484            }
 485        }
 486    }
 487
 488    fn create_markdown_block(
 489        content: String,
 490        language_registry: &Arc<LanguageRegistry>,
 491        cx: &mut App,
 492    ) -> ContentBlock {
 493        ContentBlock::Markdown {
 494            markdown: cx
 495                .new(|cx| Markdown::new(content.into(), Some(language_registry.clone()), None, cx)),
 496        }
 497    }
 498
 499    fn block_string_contents(&self, block: acp::ContentBlock) -> String {
 500        match block {
 501            acp::ContentBlock::Text(text_content) => text_content.text,
 502            acp::ContentBlock::ResourceLink(resource_link) => {
 503                Self::resource_link_md(&resource_link.uri)
 504            }
 505            acp::ContentBlock::Resource(acp::EmbeddedResource {
 506                resource:
 507                    acp::EmbeddedResourceResource::TextResourceContents(acp::TextResourceContents {
 508                        uri,
 509                        ..
 510                    }),
 511                ..
 512            }) => Self::resource_link_md(&uri),
 513            acp::ContentBlock::Image(image) => Self::image_md(&image),
 514            acp::ContentBlock::Audio(_) | acp::ContentBlock::Resource(_) => String::new(),
 515        }
 516    }
 517
 518    fn resource_link_md(uri: &str) -> String {
 519        if let Some(uri) = MentionUri::parse(uri).log_err() {
 520            uri.as_link().to_string()
 521        } else {
 522            uri.to_string()
 523        }
 524    }
 525
 526    fn image_md(_image: &acp::ImageContent) -> String {
 527        "`Image`".into()
 528    }
 529
 530    pub fn to_markdown<'a>(&'a self, cx: &'a App) -> &'a str {
 531        match self {
 532            ContentBlock::Empty => "",
 533            ContentBlock::Markdown { markdown } => markdown.read(cx).source(),
 534            ContentBlock::ResourceLink { resource_link } => &resource_link.uri,
 535        }
 536    }
 537
 538    pub fn markdown(&self) -> Option<&Entity<Markdown>> {
 539        match self {
 540            ContentBlock::Empty => None,
 541            ContentBlock::Markdown { markdown } => Some(markdown),
 542            ContentBlock::ResourceLink { .. } => None,
 543        }
 544    }
 545
 546    pub fn resource_link(&self) -> Option<&acp::ResourceLink> {
 547        match self {
 548            ContentBlock::ResourceLink { resource_link } => Some(resource_link),
 549            _ => None,
 550        }
 551    }
 552}
 553
 554#[derive(Debug)]
 555pub enum ToolCallContent {
 556    ContentBlock(ContentBlock),
 557    Diff(Entity<Diff>),
 558    Terminal(Entity<Terminal>),
 559}
 560
 561impl ToolCallContent {
 562    pub fn from_acp(
 563        content: acp::ToolCallContent,
 564        language_registry: Arc<LanguageRegistry>,
 565        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 566        cx: &mut App,
 567    ) -> Result<Self> {
 568        match content {
 569            acp::ToolCallContent::Content { content } => Ok(Self::ContentBlock(ContentBlock::new(
 570                content,
 571                &language_registry,
 572                cx,
 573            ))),
 574            acp::ToolCallContent::Diff { diff } => Ok(Self::Diff(cx.new(|cx| {
 575                Diff::finalized(
 576                    diff.path.to_string_lossy().into_owned(),
 577                    diff.old_text,
 578                    diff.new_text,
 579                    language_registry,
 580                    cx,
 581                )
 582            }))),
 583            acp::ToolCallContent::Terminal { terminal_id } => terminals
 584                .get(&terminal_id)
 585                .cloned()
 586                .map(Self::Terminal)
 587                .ok_or_else(|| anyhow::anyhow!("Terminal with id `{}` not found", terminal_id)),
 588        }
 589    }
 590
 591    pub fn update_from_acp(
 592        &mut self,
 593        new: acp::ToolCallContent,
 594        language_registry: Arc<LanguageRegistry>,
 595        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 596        cx: &mut App,
 597    ) -> Result<()> {
 598        let needs_update = match (&self, &new) {
 599            (Self::Diff(old_diff), acp::ToolCallContent::Diff { diff: new_diff }) => {
 600                old_diff.read(cx).needs_update(
 601                    new_diff.old_text.as_deref().unwrap_or(""),
 602                    &new_diff.new_text,
 603                    cx,
 604                )
 605            }
 606            _ => true,
 607        };
 608
 609        if needs_update {
 610            *self = Self::from_acp(new, language_registry, terminals, cx)?;
 611        }
 612        Ok(())
 613    }
 614
 615    pub fn to_markdown(&self, cx: &App) -> String {
 616        match self {
 617            Self::ContentBlock(content) => content.to_markdown(cx).to_string(),
 618            Self::Diff(diff) => diff.read(cx).to_markdown(cx),
 619            Self::Terminal(terminal) => terminal.read(cx).to_markdown(cx),
 620        }
 621    }
 622}
 623
 624#[derive(Debug, PartialEq)]
 625pub enum ToolCallUpdate {
 626    UpdateFields(acp::ToolCallUpdate),
 627    UpdateDiff(ToolCallUpdateDiff),
 628    UpdateTerminal(ToolCallUpdateTerminal),
 629}
 630
 631impl ToolCallUpdate {
 632    fn id(&self) -> &acp::ToolCallId {
 633        match self {
 634            Self::UpdateFields(update) => &update.id,
 635            Self::UpdateDiff(diff) => &diff.id,
 636            Self::UpdateTerminal(terminal) => &terminal.id,
 637        }
 638    }
 639}
 640
 641impl From<acp::ToolCallUpdate> for ToolCallUpdate {
 642    fn from(update: acp::ToolCallUpdate) -> Self {
 643        Self::UpdateFields(update)
 644    }
 645}
 646
 647impl From<ToolCallUpdateDiff> for ToolCallUpdate {
 648    fn from(diff: ToolCallUpdateDiff) -> Self {
 649        Self::UpdateDiff(diff)
 650    }
 651}
 652
 653#[derive(Debug, PartialEq)]
 654pub struct ToolCallUpdateDiff {
 655    pub id: acp::ToolCallId,
 656    pub diff: Entity<Diff>,
 657}
 658
 659impl From<ToolCallUpdateTerminal> for ToolCallUpdate {
 660    fn from(terminal: ToolCallUpdateTerminal) -> Self {
 661        Self::UpdateTerminal(terminal)
 662    }
 663}
 664
 665#[derive(Debug, PartialEq)]
 666pub struct ToolCallUpdateTerminal {
 667    pub id: acp::ToolCallId,
 668    pub terminal: Entity<Terminal>,
 669}
 670
 671#[derive(Debug, Default)]
 672pub struct Plan {
 673    pub entries: Vec<PlanEntry>,
 674}
 675
 676#[derive(Debug)]
 677pub struct PlanStats<'a> {
 678    pub in_progress_entry: Option<&'a PlanEntry>,
 679    pub pending: u32,
 680    pub completed: u32,
 681}
 682
 683impl Plan {
 684    pub fn is_empty(&self) -> bool {
 685        self.entries.is_empty()
 686    }
 687
 688    pub fn stats(&self) -> PlanStats<'_> {
 689        let mut stats = PlanStats {
 690            in_progress_entry: None,
 691            pending: 0,
 692            completed: 0,
 693        };
 694
 695        for entry in &self.entries {
 696            match &entry.status {
 697                acp::PlanEntryStatus::Pending => {
 698                    stats.pending += 1;
 699                }
 700                acp::PlanEntryStatus::InProgress => {
 701                    stats.in_progress_entry = stats.in_progress_entry.or(Some(entry));
 702                }
 703                acp::PlanEntryStatus::Completed => {
 704                    stats.completed += 1;
 705                }
 706            }
 707        }
 708
 709        stats
 710    }
 711}
 712
 713#[derive(Debug)]
 714pub struct PlanEntry {
 715    pub content: Entity<Markdown>,
 716    pub priority: acp::PlanEntryPriority,
 717    pub status: acp::PlanEntryStatus,
 718}
 719
 720impl PlanEntry {
 721    pub fn from_acp(entry: acp::PlanEntry, cx: &mut App) -> Self {
 722        Self {
 723            content: cx.new(|cx| Markdown::new(entry.content.into(), None, None, cx)),
 724            priority: entry.priority,
 725            status: entry.status,
 726        }
 727    }
 728}
 729
 730#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 731pub struct TokenUsage {
 732    pub max_tokens: u64,
 733    pub used_tokens: u64,
 734}
 735
 736impl TokenUsage {
 737    pub fn ratio(&self) -> TokenUsageRatio {
 738        #[cfg(debug_assertions)]
 739        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
 740            .unwrap_or("0.8".to_string())
 741            .parse()
 742            .unwrap();
 743        #[cfg(not(debug_assertions))]
 744        let warning_threshold: f32 = 0.8;
 745
 746        // When the maximum is unknown because there is no selected model,
 747        // avoid showing the token limit warning.
 748        if self.max_tokens == 0 {
 749            TokenUsageRatio::Normal
 750        } else if self.used_tokens >= self.max_tokens {
 751            TokenUsageRatio::Exceeded
 752        } else if self.used_tokens as f32 / self.max_tokens as f32 >= warning_threshold {
 753            TokenUsageRatio::Warning
 754        } else {
 755            TokenUsageRatio::Normal
 756        }
 757    }
 758}
 759
 760#[derive(Debug, Clone, PartialEq, Eq)]
 761pub enum TokenUsageRatio {
 762    Normal,
 763    Warning,
 764    Exceeded,
 765}
 766
 767#[derive(Debug, Clone)]
 768pub struct RetryStatus {
 769    pub last_error: SharedString,
 770    pub attempt: usize,
 771    pub max_attempts: usize,
 772    pub started_at: Instant,
 773    pub duration: Duration,
 774}
 775
 776pub struct AcpThread {
 777    title: SharedString,
 778    entries: Vec<AgentThreadEntry>,
 779    plan: Plan,
 780    project: Entity<Project>,
 781    action_log: Entity<ActionLog>,
 782    shared_buffers: HashMap<Entity<Buffer>, BufferSnapshot>,
 783    send_task: Option<Task<()>>,
 784    connection: Rc<dyn AgentConnection>,
 785    session_id: acp::SessionId,
 786    token_usage: Option<TokenUsage>,
 787    prompt_capabilities: acp::PromptCapabilities,
 788    _observe_prompt_capabilities: Task<anyhow::Result<()>>,
 789    terminals: HashMap<acp::TerminalId, Entity<Terminal>>,
 790}
 791
 792#[derive(Debug)]
 793pub enum AcpThreadEvent {
 794    NewEntry,
 795    TitleUpdated,
 796    TokenUsageUpdated,
 797    EntryUpdated(usize),
 798    EntriesRemoved(Range<usize>),
 799    ToolAuthorizationRequired,
 800    Retry(RetryStatus),
 801    Stopped,
 802    Error,
 803    LoadError(LoadError),
 804    PromptCapabilitiesUpdated,
 805    Refusal,
 806    AvailableCommandsUpdated(Vec<acp::AvailableCommand>),
 807    ModeUpdated(acp::SessionModeId),
 808}
 809
 810impl EventEmitter<AcpThreadEvent> for AcpThread {}
 811
 812#[derive(PartialEq, Eq, Debug)]
 813pub enum ThreadStatus {
 814    Idle,
 815    Generating,
 816}
 817
 818#[derive(Debug, Clone)]
 819pub enum LoadError {
 820    Unsupported {
 821        command: SharedString,
 822        current_version: SharedString,
 823        minimum_version: SharedString,
 824    },
 825    FailedToInstall(SharedString),
 826    Exited {
 827        status: ExitStatus,
 828    },
 829    Other(SharedString),
 830}
 831
 832impl Display for LoadError {
 833    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 834        match self {
 835            LoadError::Unsupported {
 836                command: path,
 837                current_version,
 838                minimum_version,
 839            } => {
 840                write!(
 841                    f,
 842                    "version {current_version} from {path} is not supported (need at least {minimum_version})"
 843                )
 844            }
 845            LoadError::FailedToInstall(msg) => write!(f, "Failed to install: {msg}"),
 846            LoadError::Exited { status } => write!(f, "Server exited with status {status}"),
 847            LoadError::Other(msg) => write!(f, "{msg}"),
 848        }
 849    }
 850}
 851
 852impl Error for LoadError {}
 853
 854impl AcpThread {
 855    pub fn new(
 856        title: impl Into<SharedString>,
 857        connection: Rc<dyn AgentConnection>,
 858        project: Entity<Project>,
 859        action_log: Entity<ActionLog>,
 860        session_id: acp::SessionId,
 861        mut prompt_capabilities_rx: watch::Receiver<acp::PromptCapabilities>,
 862        cx: &mut Context<Self>,
 863    ) -> Self {
 864        let prompt_capabilities = prompt_capabilities_rx.borrow().clone();
 865        let task = cx.spawn::<_, anyhow::Result<()>>(async move |this, cx| {
 866            loop {
 867                let caps = prompt_capabilities_rx.recv().await?;
 868                this.update(cx, |this, cx| {
 869                    this.prompt_capabilities = caps;
 870                    cx.emit(AcpThreadEvent::PromptCapabilitiesUpdated);
 871                })?;
 872            }
 873        });
 874
 875        Self {
 876            action_log,
 877            shared_buffers: Default::default(),
 878            entries: Default::default(),
 879            plan: Default::default(),
 880            title: title.into(),
 881            project,
 882            send_task: None,
 883            connection,
 884            session_id,
 885            token_usage: None,
 886            prompt_capabilities,
 887            _observe_prompt_capabilities: task,
 888            terminals: HashMap::default(),
 889        }
 890    }
 891
 892    pub fn prompt_capabilities(&self) -> acp::PromptCapabilities {
 893        self.prompt_capabilities.clone()
 894    }
 895
 896    pub fn connection(&self) -> &Rc<dyn AgentConnection> {
 897        &self.connection
 898    }
 899
 900    pub fn action_log(&self) -> &Entity<ActionLog> {
 901        &self.action_log
 902    }
 903
 904    pub fn project(&self) -> &Entity<Project> {
 905        &self.project
 906    }
 907
 908    pub fn title(&self) -> SharedString {
 909        self.title.clone()
 910    }
 911
 912    pub fn entries(&self) -> &[AgentThreadEntry] {
 913        &self.entries
 914    }
 915
 916    pub fn session_id(&self) -> &acp::SessionId {
 917        &self.session_id
 918    }
 919
 920    pub fn status(&self) -> ThreadStatus {
 921        if self.send_task.is_some() {
 922            ThreadStatus::Generating
 923        } else {
 924            ThreadStatus::Idle
 925        }
 926    }
 927
 928    pub fn token_usage(&self) -> Option<&TokenUsage> {
 929        self.token_usage.as_ref()
 930    }
 931
 932    pub fn has_pending_edit_tool_calls(&self) -> bool {
 933        for entry in self.entries.iter().rev() {
 934            match entry {
 935                AgentThreadEntry::UserMessage(_) => return false,
 936                AgentThreadEntry::ToolCall(
 937                    call @ ToolCall {
 938                        status: ToolCallStatus::InProgress | ToolCallStatus::Pending,
 939                        ..
 940                    },
 941                ) if call.diffs().next().is_some() => {
 942                    return true;
 943                }
 944                AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
 945            }
 946        }
 947
 948        false
 949    }
 950
 951    pub fn used_tools_since_last_user_message(&self) -> bool {
 952        for entry in self.entries.iter().rev() {
 953            match entry {
 954                AgentThreadEntry::UserMessage(..) => return false,
 955                AgentThreadEntry::AssistantMessage(..) => continue,
 956                AgentThreadEntry::ToolCall(..) => return true,
 957            }
 958        }
 959
 960        false
 961    }
 962
 963    pub fn handle_session_update(
 964        &mut self,
 965        update: acp::SessionUpdate,
 966        cx: &mut Context<Self>,
 967    ) -> Result<(), acp::Error> {
 968        match update {
 969            acp::SessionUpdate::UserMessageChunk { content } => {
 970                self.push_user_content_block(None, content, cx);
 971            }
 972            acp::SessionUpdate::AgentMessageChunk { content } => {
 973                self.push_assistant_content_block(content, false, cx);
 974            }
 975            acp::SessionUpdate::AgentThoughtChunk { content } => {
 976                self.push_assistant_content_block(content, true, cx);
 977            }
 978            acp::SessionUpdate::ToolCall(tool_call) => {
 979                self.upsert_tool_call(tool_call, cx)?;
 980            }
 981            acp::SessionUpdate::ToolCallUpdate(tool_call_update) => {
 982                self.update_tool_call(tool_call_update, cx)?;
 983            }
 984            acp::SessionUpdate::Plan(plan) => {
 985                self.update_plan(plan, cx);
 986            }
 987            acp::SessionUpdate::AvailableCommandsUpdate { available_commands } => {
 988                cx.emit(AcpThreadEvent::AvailableCommandsUpdated(available_commands))
 989            }
 990            acp::SessionUpdate::CurrentModeUpdate { current_mode_id } => {
 991                cx.emit(AcpThreadEvent::ModeUpdated(current_mode_id))
 992            }
 993        }
 994        Ok(())
 995    }
 996
 997    pub fn push_user_content_block(
 998        &mut self,
 999        message_id: Option<UserMessageId>,
1000        chunk: acp::ContentBlock,
1001        cx: &mut Context<Self>,
1002    ) {
1003        let language_registry = self.project.read(cx).languages().clone();
1004        let entries_len = self.entries.len();
1005
1006        if let Some(last_entry) = self.entries.last_mut()
1007            && let AgentThreadEntry::UserMessage(UserMessage {
1008                id,
1009                content,
1010                chunks,
1011                ..
1012            }) = last_entry
1013        {
1014            *id = message_id.or(id.take());
1015            content.append(chunk.clone(), &language_registry, cx);
1016            chunks.push(chunk);
1017            let idx = entries_len - 1;
1018            cx.emit(AcpThreadEvent::EntryUpdated(idx));
1019        } else {
1020            let content = ContentBlock::new(chunk.clone(), &language_registry, cx);
1021            self.push_entry(
1022                AgentThreadEntry::UserMessage(UserMessage {
1023                    id: message_id,
1024                    content,
1025                    chunks: vec![chunk],
1026                    checkpoint: None,
1027                }),
1028                cx,
1029            );
1030        }
1031    }
1032
1033    pub fn push_assistant_content_block(
1034        &mut self,
1035        chunk: acp::ContentBlock,
1036        is_thought: bool,
1037        cx: &mut Context<Self>,
1038    ) {
1039        let language_registry = self.project.read(cx).languages().clone();
1040        let entries_len = self.entries.len();
1041        if let Some(last_entry) = self.entries.last_mut()
1042            && let AgentThreadEntry::AssistantMessage(AssistantMessage { chunks }) = last_entry
1043        {
1044            let idx = entries_len - 1;
1045            cx.emit(AcpThreadEvent::EntryUpdated(idx));
1046            match (chunks.last_mut(), is_thought) {
1047                (Some(AssistantMessageChunk::Message { block }), false)
1048                | (Some(AssistantMessageChunk::Thought { block }), true) => {
1049                    block.append(chunk, &language_registry, cx)
1050                }
1051                _ => {
1052                    let block = ContentBlock::new(chunk, &language_registry, cx);
1053                    if is_thought {
1054                        chunks.push(AssistantMessageChunk::Thought { block })
1055                    } else {
1056                        chunks.push(AssistantMessageChunk::Message { block })
1057                    }
1058                }
1059            }
1060        } else {
1061            let block = ContentBlock::new(chunk, &language_registry, cx);
1062            let chunk = if is_thought {
1063                AssistantMessageChunk::Thought { block }
1064            } else {
1065                AssistantMessageChunk::Message { block }
1066            };
1067
1068            self.push_entry(
1069                AgentThreadEntry::AssistantMessage(AssistantMessage {
1070                    chunks: vec![chunk],
1071                }),
1072                cx,
1073            );
1074        }
1075    }
1076
1077    fn push_entry(&mut self, entry: AgentThreadEntry, cx: &mut Context<Self>) {
1078        self.entries.push(entry);
1079        cx.emit(AcpThreadEvent::NewEntry);
1080    }
1081
1082    pub fn can_set_title(&mut self, cx: &mut Context<Self>) -> bool {
1083        self.connection.set_title(&self.session_id, cx).is_some()
1084    }
1085
1086    pub fn set_title(&mut self, title: SharedString, cx: &mut Context<Self>) -> Task<Result<()>> {
1087        if title != self.title {
1088            self.title = title.clone();
1089            cx.emit(AcpThreadEvent::TitleUpdated);
1090            if let Some(set_title) = self.connection.set_title(&self.session_id, cx) {
1091                return set_title.run(title, cx);
1092            }
1093        }
1094        Task::ready(Ok(()))
1095    }
1096
1097    pub fn update_token_usage(&mut self, usage: Option<TokenUsage>, cx: &mut Context<Self>) {
1098        self.token_usage = usage;
1099        cx.emit(AcpThreadEvent::TokenUsageUpdated);
1100    }
1101
1102    pub fn update_retry_status(&mut self, status: RetryStatus, cx: &mut Context<Self>) {
1103        cx.emit(AcpThreadEvent::Retry(status));
1104    }
1105
1106    pub fn update_tool_call(
1107        &mut self,
1108        update: impl Into<ToolCallUpdate>,
1109        cx: &mut Context<Self>,
1110    ) -> Result<()> {
1111        let update = update.into();
1112        let languages = self.project.read(cx).languages().clone();
1113
1114        let ix = match self.index_for_tool_call(update.id()) {
1115            Some(ix) => ix,
1116            None => {
1117                // Tool call not found - create a failed tool call entry
1118                let failed_tool_call = ToolCall {
1119                    id: update.id().clone(),
1120                    label: cx.new(|cx| Markdown::new("Tool call not found".into(), None, None, cx)),
1121                    kind: acp::ToolKind::Fetch,
1122                    content: vec![ToolCallContent::ContentBlock(ContentBlock::new(
1123                        acp::ContentBlock::Text(acp::TextContent {
1124                            text: "Tool call not found".to_string(),
1125                            annotations: None,
1126                            meta: None,
1127                        }),
1128                        &languages,
1129                        cx,
1130                    ))],
1131                    status: ToolCallStatus::Failed,
1132                    locations: Vec::new(),
1133                    resolved_locations: Vec::new(),
1134                    raw_input: None,
1135                    raw_output: None,
1136                };
1137                self.push_entry(AgentThreadEntry::ToolCall(failed_tool_call), cx);
1138                return Ok(());
1139            }
1140        };
1141        let AgentThreadEntry::ToolCall(call) = &mut self.entries[ix] else {
1142            unreachable!()
1143        };
1144
1145        match update {
1146            ToolCallUpdate::UpdateFields(update) => {
1147                let location_updated = update.fields.locations.is_some();
1148                call.update_fields(update.fields, languages, &self.terminals, cx)?;
1149                if location_updated {
1150                    self.resolve_locations(update.id, cx);
1151                }
1152            }
1153            ToolCallUpdate::UpdateDiff(update) => {
1154                call.content.clear();
1155                call.content.push(ToolCallContent::Diff(update.diff));
1156            }
1157            ToolCallUpdate::UpdateTerminal(update) => {
1158                call.content.clear();
1159                call.content
1160                    .push(ToolCallContent::Terminal(update.terminal));
1161            }
1162        }
1163
1164        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1165
1166        Ok(())
1167    }
1168
1169    /// Updates a tool call if id matches an existing entry, otherwise inserts a new one.
1170    pub fn upsert_tool_call(
1171        &mut self,
1172        tool_call: acp::ToolCall,
1173        cx: &mut Context<Self>,
1174    ) -> Result<(), acp::Error> {
1175        let status = tool_call.status.into();
1176        self.upsert_tool_call_inner(tool_call.into(), status, cx)
1177    }
1178
1179    /// Fails if id does not match an existing entry.
1180    pub fn upsert_tool_call_inner(
1181        &mut self,
1182        update: acp::ToolCallUpdate,
1183        status: ToolCallStatus,
1184        cx: &mut Context<Self>,
1185    ) -> Result<(), acp::Error> {
1186        let language_registry = self.project.read(cx).languages().clone();
1187        let id = update.id.clone();
1188
1189        if let Some(ix) = self.index_for_tool_call(&id) {
1190            let AgentThreadEntry::ToolCall(call) = &mut self.entries[ix] else {
1191                unreachable!()
1192            };
1193
1194            call.update_fields(update.fields, language_registry, &self.terminals, cx)?;
1195            call.status = status;
1196
1197            cx.emit(AcpThreadEvent::EntryUpdated(ix));
1198        } else {
1199            let call = ToolCall::from_acp(
1200                update.try_into()?,
1201                status,
1202                language_registry,
1203                &self.terminals,
1204                cx,
1205            )?;
1206            self.push_entry(AgentThreadEntry::ToolCall(call), cx);
1207        };
1208
1209        self.resolve_locations(id, cx);
1210        Ok(())
1211    }
1212
1213    fn index_for_tool_call(&self, id: &acp::ToolCallId) -> Option<usize> {
1214        self.entries
1215            .iter()
1216            .enumerate()
1217            .rev()
1218            .find_map(|(index, entry)| {
1219                if let AgentThreadEntry::ToolCall(tool_call) = entry
1220                    && &tool_call.id == id
1221                {
1222                    Some(index)
1223                } else {
1224                    None
1225                }
1226            })
1227    }
1228
1229    fn tool_call_mut(&mut self, id: &acp::ToolCallId) -> Option<(usize, &mut ToolCall)> {
1230        // The tool call we are looking for is typically the last one, or very close to the end.
1231        // At the moment, it doesn't seem like a hashmap would be a good fit for this use case.
1232        self.entries
1233            .iter_mut()
1234            .enumerate()
1235            .rev()
1236            .find_map(|(index, tool_call)| {
1237                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1238                    && &tool_call.id == id
1239                {
1240                    Some((index, tool_call))
1241                } else {
1242                    None
1243                }
1244            })
1245    }
1246
1247    pub fn tool_call(&mut self, id: &acp::ToolCallId) -> Option<(usize, &ToolCall)> {
1248        self.entries
1249            .iter()
1250            .enumerate()
1251            .rev()
1252            .find_map(|(index, tool_call)| {
1253                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1254                    && &tool_call.id == id
1255                {
1256                    Some((index, tool_call))
1257                } else {
1258                    None
1259                }
1260            })
1261    }
1262
1263    pub fn resolve_locations(&mut self, id: acp::ToolCallId, cx: &mut Context<Self>) {
1264        let project = self.project.clone();
1265        let Some((_, tool_call)) = self.tool_call_mut(&id) else {
1266            return;
1267        };
1268        let task = tool_call.resolve_locations(project, cx);
1269        cx.spawn(async move |this, cx| {
1270            let resolved_locations = task.await;
1271            this.update(cx, |this, cx| {
1272                let project = this.project.clone();
1273                let Some((ix, tool_call)) = this.tool_call_mut(&id) else {
1274                    return;
1275                };
1276                if let Some(Some(location)) = resolved_locations.last() {
1277                    project.update(cx, |project, cx| {
1278                        if let Some(agent_location) = project.agent_location() {
1279                            let should_ignore = agent_location.buffer == location.buffer
1280                                && location
1281                                    .buffer
1282                                    .update(cx, |buffer, _| {
1283                                        let snapshot = buffer.snapshot();
1284                                        let old_position =
1285                                            agent_location.position.to_point(&snapshot);
1286                                        let new_position = location.position.to_point(&snapshot);
1287                                        // ignore this so that when we get updates from the edit tool
1288                                        // the position doesn't reset to the startof line
1289                                        old_position.row == new_position.row
1290                                            && old_position.column > new_position.column
1291                                    })
1292                                    .ok()
1293                                    .unwrap_or_default();
1294                            if !should_ignore {
1295                                project.set_agent_location(Some(location.clone()), cx);
1296                            }
1297                        }
1298                    });
1299                }
1300                if tool_call.resolved_locations != resolved_locations {
1301                    tool_call.resolved_locations = resolved_locations;
1302                    cx.emit(AcpThreadEvent::EntryUpdated(ix));
1303                }
1304            })
1305        })
1306        .detach();
1307    }
1308
1309    pub fn request_tool_call_authorization(
1310        &mut self,
1311        tool_call: acp::ToolCallUpdate,
1312        options: Vec<acp::PermissionOption>,
1313        respect_always_allow_setting: bool,
1314        cx: &mut Context<Self>,
1315    ) -> Result<BoxFuture<'static, acp::RequestPermissionOutcome>> {
1316        let (tx, rx) = oneshot::channel();
1317
1318        if respect_always_allow_setting && AgentSettings::get_global(cx).always_allow_tool_actions {
1319            // Don't use AllowAlways, because then if you were to turn off always_allow_tool_actions,
1320            // some tools would (incorrectly) continue to auto-accept.
1321            if let Some(allow_once_option) = options.iter().find_map(|option| {
1322                if matches!(option.kind, acp::PermissionOptionKind::AllowOnce) {
1323                    Some(option.id.clone())
1324                } else {
1325                    None
1326                }
1327            }) {
1328                self.upsert_tool_call_inner(tool_call, ToolCallStatus::Pending, cx)?;
1329                return Ok(async {
1330                    acp::RequestPermissionOutcome::Selected {
1331                        option_id: allow_once_option,
1332                    }
1333                }
1334                .boxed());
1335            }
1336        }
1337
1338        let status = ToolCallStatus::WaitingForConfirmation {
1339            options,
1340            respond_tx: tx,
1341        };
1342
1343        self.upsert_tool_call_inner(tool_call, status, cx)?;
1344        cx.emit(AcpThreadEvent::ToolAuthorizationRequired);
1345
1346        let fut = async {
1347            match rx.await {
1348                Ok(option) => acp::RequestPermissionOutcome::Selected { option_id: option },
1349                Err(oneshot::Canceled) => acp::RequestPermissionOutcome::Cancelled,
1350            }
1351        }
1352        .boxed();
1353
1354        Ok(fut)
1355    }
1356
1357    pub fn authorize_tool_call(
1358        &mut self,
1359        id: acp::ToolCallId,
1360        option_id: acp::PermissionOptionId,
1361        option_kind: acp::PermissionOptionKind,
1362        cx: &mut Context<Self>,
1363    ) {
1364        let Some((ix, call)) = self.tool_call_mut(&id) else {
1365            return;
1366        };
1367
1368        let new_status = match option_kind {
1369            acp::PermissionOptionKind::RejectOnce | acp::PermissionOptionKind::RejectAlways => {
1370                ToolCallStatus::Rejected
1371            }
1372            acp::PermissionOptionKind::AllowOnce | acp::PermissionOptionKind::AllowAlways => {
1373                ToolCallStatus::InProgress
1374            }
1375        };
1376
1377        let curr_status = mem::replace(&mut call.status, new_status);
1378
1379        if let ToolCallStatus::WaitingForConfirmation { respond_tx, .. } = curr_status {
1380            respond_tx.send(option_id).log_err();
1381        } else if cfg!(debug_assertions) {
1382            panic!("tried to authorize an already authorized tool call");
1383        }
1384
1385        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1386    }
1387
1388    pub fn first_tool_awaiting_confirmation(&self) -> Option<&ToolCall> {
1389        let mut first_tool_call = None;
1390
1391        for entry in self.entries.iter().rev() {
1392            match &entry {
1393                AgentThreadEntry::ToolCall(call) => {
1394                    if let ToolCallStatus::WaitingForConfirmation { .. } = call.status {
1395                        first_tool_call = Some(call);
1396                    } else {
1397                        continue;
1398                    }
1399                }
1400                AgentThreadEntry::UserMessage(_) | AgentThreadEntry::AssistantMessage(_) => {
1401                    // Reached the beginning of the turn.
1402                    // If we had pending permission requests in the previous turn, they have been cancelled.
1403                    break;
1404                }
1405            }
1406        }
1407
1408        first_tool_call
1409    }
1410
1411    pub fn plan(&self) -> &Plan {
1412        &self.plan
1413    }
1414
1415    pub fn update_plan(&mut self, request: acp::Plan, cx: &mut Context<Self>) {
1416        let new_entries_len = request.entries.len();
1417        let mut new_entries = request.entries.into_iter();
1418
1419        // Reuse existing markdown to prevent flickering
1420        for (old, new) in self.plan.entries.iter_mut().zip(new_entries.by_ref()) {
1421            let PlanEntry {
1422                content,
1423                priority,
1424                status,
1425            } = old;
1426            content.update(cx, |old, cx| {
1427                old.replace(new.content, cx);
1428            });
1429            *priority = new.priority;
1430            *status = new.status;
1431        }
1432        for new in new_entries {
1433            self.plan.entries.push(PlanEntry::from_acp(new, cx))
1434        }
1435        self.plan.entries.truncate(new_entries_len);
1436
1437        cx.notify();
1438    }
1439
1440    fn clear_completed_plan_entries(&mut self, cx: &mut Context<Self>) {
1441        self.plan
1442            .entries
1443            .retain(|entry| !matches!(entry.status, acp::PlanEntryStatus::Completed));
1444        cx.notify();
1445    }
1446
1447    #[cfg(any(test, feature = "test-support"))]
1448    pub fn send_raw(
1449        &mut self,
1450        message: &str,
1451        cx: &mut Context<Self>,
1452    ) -> BoxFuture<'static, Result<()>> {
1453        self.send(
1454            vec![acp::ContentBlock::Text(acp::TextContent {
1455                text: message.to_string(),
1456                annotations: None,
1457                meta: None,
1458            })],
1459            cx,
1460        )
1461    }
1462
1463    pub fn send(
1464        &mut self,
1465        message: Vec<acp::ContentBlock>,
1466        cx: &mut Context<Self>,
1467    ) -> BoxFuture<'static, Result<()>> {
1468        let block = ContentBlock::new_combined(
1469            message.clone(),
1470            self.project.read(cx).languages().clone(),
1471            cx,
1472        );
1473        let request = acp::PromptRequest {
1474            prompt: message.clone(),
1475            session_id: self.session_id.clone(),
1476            meta: None,
1477        };
1478        let git_store = self.project.read(cx).git_store().clone();
1479
1480        let message_id = if self.connection.truncate(&self.session_id, cx).is_some() {
1481            Some(UserMessageId::new())
1482        } else {
1483            None
1484        };
1485
1486        self.run_turn(cx, async move |this, cx| {
1487            this.update(cx, |this, cx| {
1488                this.push_entry(
1489                    AgentThreadEntry::UserMessage(UserMessage {
1490                        id: message_id.clone(),
1491                        content: block,
1492                        chunks: message,
1493                        checkpoint: None,
1494                    }),
1495                    cx,
1496                );
1497            })
1498            .ok();
1499
1500            let old_checkpoint = git_store
1501                .update(cx, |git, cx| git.checkpoint(cx))?
1502                .await
1503                .context("failed to get old checkpoint")
1504                .log_err();
1505            this.update(cx, |this, cx| {
1506                if let Some((_ix, message)) = this.last_user_message() {
1507                    message.checkpoint = old_checkpoint.map(|git_checkpoint| Checkpoint {
1508                        git_checkpoint,
1509                        show: false,
1510                    });
1511                }
1512                this.connection.prompt(message_id, request, cx)
1513            })?
1514            .await
1515        })
1516    }
1517
1518    pub fn can_resume(&self, cx: &App) -> bool {
1519        self.connection.resume(&self.session_id, cx).is_some()
1520    }
1521
1522    pub fn resume(&mut self, cx: &mut Context<Self>) -> BoxFuture<'static, Result<()>> {
1523        self.run_turn(cx, async move |this, cx| {
1524            this.update(cx, |this, cx| {
1525                this.connection
1526                    .resume(&this.session_id, cx)
1527                    .map(|resume| resume.run(cx))
1528            })?
1529            .context("resuming a session is not supported")?
1530            .await
1531        })
1532    }
1533
1534    fn run_turn(
1535        &mut self,
1536        cx: &mut Context<Self>,
1537        f: impl 'static + AsyncFnOnce(WeakEntity<Self>, &mut AsyncApp) -> Result<acp::PromptResponse>,
1538    ) -> BoxFuture<'static, Result<()>> {
1539        self.clear_completed_plan_entries(cx);
1540
1541        let (tx, rx) = oneshot::channel();
1542        let cancel_task = self.cancel(cx);
1543
1544        self.send_task = Some(cx.spawn(async move |this, cx| {
1545            cancel_task.await;
1546            tx.send(f(this, cx).await).ok();
1547        }));
1548
1549        cx.spawn(async move |this, cx| {
1550            let response = rx.await;
1551
1552            this.update(cx, |this, cx| this.update_last_checkpoint(cx))?
1553                .await?;
1554
1555            this.update(cx, |this, cx| {
1556                this.project
1557                    .update(cx, |project, cx| project.set_agent_location(None, cx));
1558                match response {
1559                    Ok(Err(e)) => {
1560                        this.send_task.take();
1561                        cx.emit(AcpThreadEvent::Error);
1562                        Err(e)
1563                    }
1564                    result => {
1565                        let canceled = matches!(
1566                            result,
1567                            Ok(Ok(acp::PromptResponse {
1568                                stop_reason: acp::StopReason::Cancelled,
1569                                meta: None,
1570                            }))
1571                        );
1572
1573                        // We only take the task if the current prompt wasn't canceled.
1574                        //
1575                        // This prompt may have been canceled because another one was sent
1576                        // while it was still generating. In these cases, dropping `send_task`
1577                        // would cause the next generation to be canceled.
1578                        if !canceled {
1579                            this.send_task.take();
1580                        }
1581
1582                        // Handle refusal - distinguish between user prompt and tool call refusals
1583                        if let Ok(Ok(acp::PromptResponse {
1584                            stop_reason: acp::StopReason::Refusal,
1585                            meta: _,
1586                        })) = result
1587                        {
1588                            if let Some((user_msg_ix, _)) = this.last_user_message() {
1589                                // Check if there's a completed tool call with results after the last user message
1590                                // This indicates the refusal is in response to tool output, not the user's prompt
1591                                let has_completed_tool_call_after_user_msg =
1592                                    this.entries.iter().skip(user_msg_ix + 1).any(|entry| {
1593                                        if let AgentThreadEntry::ToolCall(tool_call) = entry {
1594                                            // Check if the tool call has completed and has output
1595                                            matches!(tool_call.status, ToolCallStatus::Completed)
1596                                                && tool_call.raw_output.is_some()
1597                                        } else {
1598                                            false
1599                                        }
1600                                    });
1601
1602                                if has_completed_tool_call_after_user_msg {
1603                                    // Refusal is due to tool output - don't truncate, just notify
1604                                    // The model refused based on what the tool returned
1605                                    cx.emit(AcpThreadEvent::Refusal);
1606                                } else {
1607                                    // User prompt was refused - truncate back to before the user message
1608                                    let range = user_msg_ix..this.entries.len();
1609                                    if range.start < range.end {
1610                                        this.entries.truncate(user_msg_ix);
1611                                        cx.emit(AcpThreadEvent::EntriesRemoved(range));
1612                                    }
1613                                    cx.emit(AcpThreadEvent::Refusal);
1614                                }
1615                            } else {
1616                                // No user message found, treat as general refusal
1617                                cx.emit(AcpThreadEvent::Refusal);
1618                            }
1619                        }
1620
1621                        cx.emit(AcpThreadEvent::Stopped);
1622                        Ok(())
1623                    }
1624                }
1625            })?
1626        })
1627        .boxed()
1628    }
1629
1630    pub fn cancel(&mut self, cx: &mut Context<Self>) -> Task<()> {
1631        let Some(send_task) = self.send_task.take() else {
1632            return Task::ready(());
1633        };
1634
1635        for entry in self.entries.iter_mut() {
1636            if let AgentThreadEntry::ToolCall(call) = entry {
1637                let cancel = matches!(
1638                    call.status,
1639                    ToolCallStatus::Pending
1640                        | ToolCallStatus::WaitingForConfirmation { .. }
1641                        | ToolCallStatus::InProgress
1642                );
1643
1644                if cancel {
1645                    call.status = ToolCallStatus::Canceled;
1646                }
1647            }
1648        }
1649
1650        self.connection.cancel(&self.session_id, cx);
1651
1652        // Wait for the send task to complete
1653        cx.foreground_executor().spawn(send_task)
1654    }
1655
1656    /// Restores the git working tree to the state at the given checkpoint (if one exists)
1657    pub fn restore_checkpoint(
1658        &mut self,
1659        id: UserMessageId,
1660        cx: &mut Context<Self>,
1661    ) -> Task<Result<()>> {
1662        let Some((_, message)) = self.user_message_mut(&id) else {
1663            return Task::ready(Err(anyhow!("message not found")));
1664        };
1665
1666        let checkpoint = message
1667            .checkpoint
1668            .as_ref()
1669            .map(|c| c.git_checkpoint.clone());
1670        let rewind = self.rewind(id.clone(), cx);
1671        let git_store = self.project.read(cx).git_store().clone();
1672
1673        cx.spawn(async move |_, cx| {
1674            rewind.await?;
1675            if let Some(checkpoint) = checkpoint {
1676                git_store
1677                    .update(cx, |git, cx| git.restore_checkpoint(checkpoint, cx))?
1678                    .await?;
1679            }
1680
1681            Ok(())
1682        })
1683    }
1684
1685    /// Rewinds this thread to before the entry at `index`, removing it and all
1686    /// subsequent entries while rejecting any action_log changes made from that point.
1687    /// Unlike `restore_checkpoint`, this method does not restore from git.
1688    pub fn rewind(&mut self, id: UserMessageId, cx: &mut Context<Self>) -> Task<Result<()>> {
1689        let Some(truncate) = self.connection.truncate(&self.session_id, cx) else {
1690            return Task::ready(Err(anyhow!("not supported")));
1691        };
1692
1693        cx.spawn(async move |this, cx| {
1694            cx.update(|cx| truncate.run(id.clone(), cx))?.await?;
1695            this.update(cx, |this, cx| {
1696                if let Some((ix, _)) = this.user_message_mut(&id) {
1697                    let range = ix..this.entries.len();
1698                    this.entries.truncate(ix);
1699                    cx.emit(AcpThreadEvent::EntriesRemoved(range));
1700                }
1701                this.action_log()
1702                    .update(cx, |action_log, cx| action_log.reject_all_edits(cx))
1703            })?
1704            .await;
1705            Ok(())
1706        })
1707    }
1708
1709    fn update_last_checkpoint(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
1710        let git_store = self.project.read(cx).git_store().clone();
1711
1712        let old_checkpoint = if let Some((_, message)) = self.last_user_message() {
1713            if let Some(checkpoint) = message.checkpoint.as_ref() {
1714                checkpoint.git_checkpoint.clone()
1715            } else {
1716                return Task::ready(Ok(()));
1717            }
1718        } else {
1719            return Task::ready(Ok(()));
1720        };
1721
1722        let new_checkpoint = git_store.update(cx, |git, cx| git.checkpoint(cx));
1723        cx.spawn(async move |this, cx| {
1724            let new_checkpoint = new_checkpoint
1725                .await
1726                .context("failed to get new checkpoint")
1727                .log_err();
1728            if let Some(new_checkpoint) = new_checkpoint {
1729                let equal = git_store
1730                    .update(cx, |git, cx| {
1731                        git.compare_checkpoints(old_checkpoint.clone(), new_checkpoint, cx)
1732                    })?
1733                    .await
1734                    .unwrap_or(true);
1735                this.update(cx, |this, cx| {
1736                    let (ix, message) = this.last_user_message().context("no user message")?;
1737                    let checkpoint = message.checkpoint.as_mut().context("no checkpoint")?;
1738                    checkpoint.show = !equal;
1739                    cx.emit(AcpThreadEvent::EntryUpdated(ix));
1740                    anyhow::Ok(())
1741                })??;
1742            }
1743
1744            Ok(())
1745        })
1746    }
1747
1748    fn last_user_message(&mut self) -> Option<(usize, &mut UserMessage)> {
1749        self.entries
1750            .iter_mut()
1751            .enumerate()
1752            .rev()
1753            .find_map(|(ix, entry)| {
1754                if let AgentThreadEntry::UserMessage(message) = entry {
1755                    Some((ix, message))
1756                } else {
1757                    None
1758                }
1759            })
1760    }
1761
1762    fn user_message_mut(&mut self, id: &UserMessageId) -> Option<(usize, &mut UserMessage)> {
1763        self.entries.iter_mut().enumerate().find_map(|(ix, entry)| {
1764            if let AgentThreadEntry::UserMessage(message) = entry {
1765                if message.id.as_ref() == Some(id) {
1766                    Some((ix, message))
1767                } else {
1768                    None
1769                }
1770            } else {
1771                None
1772            }
1773        })
1774    }
1775
1776    pub fn read_text_file(
1777        &self,
1778        path: PathBuf,
1779        line: Option<u32>,
1780        limit: Option<u32>,
1781        reuse_shared_snapshot: bool,
1782        cx: &mut Context<Self>,
1783    ) -> Task<Result<String, acp::Error>> {
1784        // Args are 1-based, move to 0-based
1785        let line = line.unwrap_or_default().saturating_sub(1);
1786        let limit = limit.unwrap_or(u32::MAX);
1787        let project = self.project.clone();
1788        let action_log = self.action_log.clone();
1789        cx.spawn(async move |this, cx| {
1790            let load = project
1791                .update(cx, |project, cx| {
1792                    let path = project
1793                        .project_path_for_absolute_path(&path, cx)
1794                        .ok_or_else(|| {
1795                            acp::Error::resource_not_found(Some(path.display().to_string()))
1796                        })?;
1797                    Ok(project.open_buffer(path, cx))
1798                })
1799                .map_err(|e| acp::Error::internal_error().with_data(e.to_string()))
1800                .flatten()?;
1801
1802            let buffer = load.await?;
1803
1804            let snapshot = if reuse_shared_snapshot {
1805                this.read_with(cx, |this, _| {
1806                    this.shared_buffers.get(&buffer.clone()).cloned()
1807                })
1808                .log_err()
1809                .flatten()
1810            } else {
1811                None
1812            };
1813
1814            let snapshot = if let Some(snapshot) = snapshot {
1815                snapshot
1816            } else {
1817                action_log.update(cx, |action_log, cx| {
1818                    action_log.buffer_read(buffer.clone(), cx);
1819                })?;
1820
1821                let snapshot = buffer.update(cx, |buffer, _| buffer.snapshot())?;
1822                this.update(cx, |this, _| {
1823                    this.shared_buffers.insert(buffer.clone(), snapshot.clone());
1824                })?;
1825                snapshot
1826            };
1827
1828            let max_point = snapshot.max_point();
1829            let start_position = Point::new(line, 0);
1830
1831            if start_position > max_point {
1832                return Err(acp::Error::invalid_params().with_data(format!(
1833                    "Attempting to read beyond the end of the file, line {}:{}",
1834                    max_point.row + 1,
1835                    max_point.column
1836                )));
1837            }
1838
1839            let start = snapshot.anchor_before(start_position);
1840            let end = snapshot.anchor_before(Point::new(line.saturating_add(limit), 0));
1841
1842            project.update(cx, |project, cx| {
1843                project.set_agent_location(
1844                    Some(AgentLocation {
1845                        buffer: buffer.downgrade(),
1846                        position: start,
1847                    }),
1848                    cx,
1849                );
1850            })?;
1851
1852            Ok(snapshot.text_for_range(start..end).collect::<String>())
1853        })
1854    }
1855
1856    pub fn write_text_file(
1857        &self,
1858        path: PathBuf,
1859        content: String,
1860        cx: &mut Context<Self>,
1861    ) -> Task<Result<()>> {
1862        let project = self.project.clone();
1863        let action_log = self.action_log.clone();
1864        cx.spawn(async move |this, cx| {
1865            let load = project.update(cx, |project, cx| {
1866                let path = project
1867                    .project_path_for_absolute_path(&path, cx)
1868                    .context("invalid path")?;
1869                anyhow::Ok(project.open_buffer(path, cx))
1870            });
1871            let buffer = load??.await?;
1872            let snapshot = this.update(cx, |this, cx| {
1873                this.shared_buffers
1874                    .get(&buffer)
1875                    .cloned()
1876                    .unwrap_or_else(|| buffer.read(cx).snapshot())
1877            })?;
1878            let edits = cx
1879                .background_executor()
1880                .spawn(async move {
1881                    let old_text = snapshot.text();
1882                    text_diff(old_text.as_str(), &content)
1883                        .into_iter()
1884                        .map(|(range, replacement)| {
1885                            (
1886                                snapshot.anchor_after(range.start)
1887                                    ..snapshot.anchor_before(range.end),
1888                                replacement,
1889                            )
1890                        })
1891                        .collect::<Vec<_>>()
1892                })
1893                .await;
1894
1895            project.update(cx, |project, cx| {
1896                project.set_agent_location(
1897                    Some(AgentLocation {
1898                        buffer: buffer.downgrade(),
1899                        position: edits
1900                            .last()
1901                            .map(|(range, _)| range.end)
1902                            .unwrap_or(Anchor::MIN),
1903                    }),
1904                    cx,
1905                );
1906            })?;
1907
1908            let format_on_save = cx.update(|cx| {
1909                action_log.update(cx, |action_log, cx| {
1910                    action_log.buffer_read(buffer.clone(), cx);
1911                });
1912
1913                let format_on_save = buffer.update(cx, |buffer, cx| {
1914                    buffer.edit(edits, None, cx);
1915
1916                    let settings = language::language_settings::language_settings(
1917                        buffer.language().map(|l| l.name()),
1918                        buffer.file(),
1919                        cx,
1920                    );
1921
1922                    settings.format_on_save != FormatOnSave::Off
1923                });
1924                action_log.update(cx, |action_log, cx| {
1925                    action_log.buffer_edited(buffer.clone(), cx);
1926                });
1927                format_on_save
1928            })?;
1929
1930            if format_on_save {
1931                let format_task = project.update(cx, |project, cx| {
1932                    project.format(
1933                        HashSet::from_iter([buffer.clone()]),
1934                        LspFormatTarget::Buffers,
1935                        false,
1936                        FormatTrigger::Save,
1937                        cx,
1938                    )
1939                })?;
1940                format_task.await.log_err();
1941
1942                action_log.update(cx, |action_log, cx| {
1943                    action_log.buffer_edited(buffer.clone(), cx);
1944                })?;
1945            }
1946
1947            project
1948                .update(cx, |project, cx| project.save_buffer(buffer, cx))?
1949                .await
1950        })
1951    }
1952
1953    pub fn create_terminal(
1954        &self,
1955        command: String,
1956        args: Vec<String>,
1957        extra_env: Vec<acp::EnvVariable>,
1958        cwd: Option<PathBuf>,
1959        output_byte_limit: Option<u64>,
1960        cx: &mut Context<Self>,
1961    ) -> Task<Result<Entity<Terminal>>> {
1962        let env = match &cwd {
1963            Some(dir) => self.project.update(cx, |project, cx| {
1964                project.directory_environment(dir.as_path().into(), cx)
1965            }),
1966            None => Task::ready(None).shared(),
1967        };
1968
1969        let env = cx.spawn(async move |_, _| {
1970            let mut env = env.await.unwrap_or_default();
1971            if cfg!(unix) {
1972                env.insert("PAGER".into(), "cat".into());
1973            }
1974            for var in extra_env {
1975                env.insert(var.name, var.value);
1976            }
1977            env
1978        });
1979
1980        let project = self.project.clone();
1981        let language_registry = project.read(cx).languages().clone();
1982
1983        let terminal_id = acp::TerminalId(Uuid::new_v4().to_string().into());
1984        let terminal_task = cx.spawn({
1985            let terminal_id = terminal_id.clone();
1986            async move |_this, cx| {
1987                let env = env.await;
1988                let (task_command, task_args) = ShellBuilder::new(
1989                    project
1990                        .update(cx, |project, cx| {
1991                            project
1992                                .remote_client()
1993                                .and_then(|r| r.read(cx).default_system_shell())
1994                        })?
1995                        .as_deref(),
1996                    &Shell::Program(get_default_system_shell()),
1997                )
1998                .redirect_stdin_to_dev_null()
1999                .build(Some(command.clone()), &args);
2000                let terminal = project
2001                    .update(cx, |project, cx| {
2002                        project.create_terminal_task(
2003                            task::SpawnInTerminal {
2004                                command: Some(task_command),
2005                                args: task_args,
2006                                cwd: cwd.clone(),
2007                                env,
2008                                ..Default::default()
2009                            },
2010                            cx,
2011                        )
2012                    })?
2013                    .await?;
2014
2015                cx.new(|cx| {
2016                    Terminal::new(
2017                        terminal_id,
2018                        &format!("{} {}", command, args.join(" ")),
2019                        cwd,
2020                        output_byte_limit.map(|l| l as usize),
2021                        terminal,
2022                        language_registry,
2023                        cx,
2024                    )
2025                })
2026            }
2027        });
2028
2029        cx.spawn(async move |this, cx| {
2030            let terminal = terminal_task.await?;
2031            this.update(cx, |this, _cx| {
2032                this.terminals.insert(terminal_id, terminal.clone());
2033                terminal
2034            })
2035        })
2036    }
2037
2038    pub fn kill_terminal(
2039        &mut self,
2040        terminal_id: acp::TerminalId,
2041        cx: &mut Context<Self>,
2042    ) -> Result<()> {
2043        self.terminals
2044            .get(&terminal_id)
2045            .context("Terminal not found")?
2046            .update(cx, |terminal, cx| {
2047                terminal.kill(cx);
2048            });
2049
2050        Ok(())
2051    }
2052
2053    pub fn release_terminal(
2054        &mut self,
2055        terminal_id: acp::TerminalId,
2056        cx: &mut Context<Self>,
2057    ) -> Result<()> {
2058        self.terminals
2059            .remove(&terminal_id)
2060            .context("Terminal not found")?
2061            .update(cx, |terminal, cx| {
2062                terminal.kill(cx);
2063            });
2064
2065        Ok(())
2066    }
2067
2068    pub fn terminal(&self, terminal_id: acp::TerminalId) -> Result<Entity<Terminal>> {
2069        self.terminals
2070            .get(&terminal_id)
2071            .context("Terminal not found")
2072            .cloned()
2073    }
2074
2075    pub fn to_markdown(&self, cx: &App) -> String {
2076        self.entries.iter().map(|e| e.to_markdown(cx)).collect()
2077    }
2078
2079    pub fn emit_load_error(&mut self, error: LoadError, cx: &mut Context<Self>) {
2080        cx.emit(AcpThreadEvent::LoadError(error));
2081    }
2082}
2083
2084fn markdown_for_raw_output(
2085    raw_output: &serde_json::Value,
2086    language_registry: &Arc<LanguageRegistry>,
2087    cx: &mut App,
2088) -> Option<Entity<Markdown>> {
2089    match raw_output {
2090        serde_json::Value::Null => None,
2091        serde_json::Value::Bool(value) => Some(cx.new(|cx| {
2092            Markdown::new(
2093                value.to_string().into(),
2094                Some(language_registry.clone()),
2095                None,
2096                cx,
2097            )
2098        })),
2099        serde_json::Value::Number(value) => Some(cx.new(|cx| {
2100            Markdown::new(
2101                value.to_string().into(),
2102                Some(language_registry.clone()),
2103                None,
2104                cx,
2105            )
2106        })),
2107        serde_json::Value::String(value) => Some(cx.new(|cx| {
2108            Markdown::new(
2109                value.clone().into(),
2110                Some(language_registry.clone()),
2111                None,
2112                cx,
2113            )
2114        })),
2115        value => Some(cx.new(|cx| {
2116            Markdown::new(
2117                format!("```json\n{}\n```", value).into(),
2118                Some(language_registry.clone()),
2119                None,
2120                cx,
2121            )
2122        })),
2123    }
2124}
2125
2126#[cfg(test)]
2127mod tests {
2128    use super::*;
2129    use anyhow::anyhow;
2130    use futures::{channel::mpsc, future::LocalBoxFuture, select};
2131    use gpui::{App, AsyncApp, TestAppContext, WeakEntity};
2132    use indoc::indoc;
2133    use project::{FakeFs, Fs};
2134    use rand::{distr, prelude::*};
2135    use serde_json::json;
2136    use settings::SettingsStore;
2137    use smol::stream::StreamExt as _;
2138    use std::{
2139        any::Any,
2140        cell::RefCell,
2141        path::Path,
2142        rc::Rc,
2143        sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
2144        time::Duration,
2145    };
2146    use util::path;
2147
2148    fn init_test(cx: &mut TestAppContext) {
2149        env_logger::try_init().ok();
2150        cx.update(|cx| {
2151            let settings_store = SettingsStore::test(cx);
2152            cx.set_global(settings_store);
2153            Project::init_settings(cx);
2154            language::init(cx);
2155        });
2156    }
2157
2158    #[gpui::test]
2159    async fn test_push_user_content_block(cx: &mut gpui::TestAppContext) {
2160        init_test(cx);
2161
2162        let fs = FakeFs::new(cx.executor());
2163        let project = Project::test(fs, [], cx).await;
2164        let connection = Rc::new(FakeAgentConnection::new());
2165        let thread = cx
2166            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2167            .await
2168            .unwrap();
2169
2170        // Test creating a new user message
2171        thread.update(cx, |thread, cx| {
2172            thread.push_user_content_block(
2173                None,
2174                acp::ContentBlock::Text(acp::TextContent {
2175                    annotations: None,
2176                    text: "Hello, ".to_string(),
2177                    meta: None,
2178                }),
2179                cx,
2180            );
2181        });
2182
2183        thread.update(cx, |thread, cx| {
2184            assert_eq!(thread.entries.len(), 1);
2185            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
2186                assert_eq!(user_msg.id, None);
2187                assert_eq!(user_msg.content.to_markdown(cx), "Hello, ");
2188            } else {
2189                panic!("Expected UserMessage");
2190            }
2191        });
2192
2193        // Test appending to existing user message
2194        let message_1_id = UserMessageId::new();
2195        thread.update(cx, |thread, cx| {
2196            thread.push_user_content_block(
2197                Some(message_1_id.clone()),
2198                acp::ContentBlock::Text(acp::TextContent {
2199                    annotations: None,
2200                    text: "world!".to_string(),
2201                    meta: None,
2202                }),
2203                cx,
2204            );
2205        });
2206
2207        thread.update(cx, |thread, cx| {
2208            assert_eq!(thread.entries.len(), 1);
2209            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
2210                assert_eq!(user_msg.id, Some(message_1_id));
2211                assert_eq!(user_msg.content.to_markdown(cx), "Hello, world!");
2212            } else {
2213                panic!("Expected UserMessage");
2214            }
2215        });
2216
2217        // Test creating new user message after assistant message
2218        thread.update(cx, |thread, cx| {
2219            thread.push_assistant_content_block(
2220                acp::ContentBlock::Text(acp::TextContent {
2221                    annotations: None,
2222                    text: "Assistant response".to_string(),
2223                    meta: None,
2224                }),
2225                false,
2226                cx,
2227            );
2228        });
2229
2230        let message_2_id = UserMessageId::new();
2231        thread.update(cx, |thread, cx| {
2232            thread.push_user_content_block(
2233                Some(message_2_id.clone()),
2234                acp::ContentBlock::Text(acp::TextContent {
2235                    annotations: None,
2236                    text: "New user message".to_string(),
2237                    meta: None,
2238                }),
2239                cx,
2240            );
2241        });
2242
2243        thread.update(cx, |thread, cx| {
2244            assert_eq!(thread.entries.len(), 3);
2245            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[2] {
2246                assert_eq!(user_msg.id, Some(message_2_id));
2247                assert_eq!(user_msg.content.to_markdown(cx), "New user message");
2248            } else {
2249                panic!("Expected UserMessage at index 2");
2250            }
2251        });
2252    }
2253
2254    #[gpui::test]
2255    async fn test_thinking_concatenation(cx: &mut gpui::TestAppContext) {
2256        init_test(cx);
2257
2258        let fs = FakeFs::new(cx.executor());
2259        let project = Project::test(fs, [], cx).await;
2260        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
2261            |_, thread, mut cx| {
2262                async move {
2263                    thread.update(&mut cx, |thread, cx| {
2264                        thread
2265                            .handle_session_update(
2266                                acp::SessionUpdate::AgentThoughtChunk {
2267                                    content: "Thinking ".into(),
2268                                },
2269                                cx,
2270                            )
2271                            .unwrap();
2272                        thread
2273                            .handle_session_update(
2274                                acp::SessionUpdate::AgentThoughtChunk {
2275                                    content: "hard!".into(),
2276                                },
2277                                cx,
2278                            )
2279                            .unwrap();
2280                    })?;
2281                    Ok(acp::PromptResponse {
2282                        stop_reason: acp::StopReason::EndTurn,
2283                        meta: None,
2284                    })
2285                }
2286                .boxed_local()
2287            },
2288        ));
2289
2290        let thread = cx
2291            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2292            .await
2293            .unwrap();
2294
2295        thread
2296            .update(cx, |thread, cx| thread.send_raw("Hello from Zed!", cx))
2297            .await
2298            .unwrap();
2299
2300        let output = thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
2301        assert_eq!(
2302            output,
2303            indoc! {r#"
2304            ## User
2305
2306            Hello from Zed!
2307
2308            ## Assistant
2309
2310            <thinking>
2311            Thinking hard!
2312            </thinking>
2313
2314            "#}
2315        );
2316    }
2317
2318    #[gpui::test]
2319    async fn test_edits_concurrently_to_user(cx: &mut TestAppContext) {
2320        init_test(cx);
2321
2322        let fs = FakeFs::new(cx.executor());
2323        fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\n"}))
2324            .await;
2325        let project = Project::test(fs.clone(), [], cx).await;
2326        let (read_file_tx, read_file_rx) = oneshot::channel::<()>();
2327        let read_file_tx = Rc::new(RefCell::new(Some(read_file_tx)));
2328        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
2329            move |_, thread, mut cx| {
2330                let read_file_tx = read_file_tx.clone();
2331                async move {
2332                    let content = thread
2333                        .update(&mut cx, |thread, cx| {
2334                            thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
2335                        })
2336                        .unwrap()
2337                        .await
2338                        .unwrap();
2339                    assert_eq!(content, "one\ntwo\nthree\n");
2340                    read_file_tx.take().unwrap().send(()).unwrap();
2341                    thread
2342                        .update(&mut cx, |thread, cx| {
2343                            thread.write_text_file(
2344                                path!("/tmp/foo").into(),
2345                                "one\ntwo\nthree\nfour\nfive\n".to_string(),
2346                                cx,
2347                            )
2348                        })
2349                        .unwrap()
2350                        .await
2351                        .unwrap();
2352                    Ok(acp::PromptResponse {
2353                        stop_reason: acp::StopReason::EndTurn,
2354                        meta: None,
2355                    })
2356                }
2357                .boxed_local()
2358            },
2359        ));
2360
2361        let (worktree, pathbuf) = project
2362            .update(cx, |project, cx| {
2363                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
2364            })
2365            .await
2366            .unwrap();
2367        let buffer = project
2368            .update(cx, |project, cx| {
2369                project.open_buffer((worktree.read(cx).id(), pathbuf), cx)
2370            })
2371            .await
2372            .unwrap();
2373
2374        let thread = cx
2375            .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2376            .await
2377            .unwrap();
2378
2379        let request = thread.update(cx, |thread, cx| {
2380            thread.send_raw("Extend the count in /tmp/foo", cx)
2381        });
2382        read_file_rx.await.ok();
2383        buffer.update(cx, |buffer, cx| {
2384            buffer.edit([(0..0, "zero\n".to_string())], None, cx);
2385        });
2386        cx.run_until_parked();
2387        assert_eq!(
2388            buffer.read_with(cx, |buffer, _| buffer.text()),
2389            "zero\none\ntwo\nthree\nfour\nfive\n"
2390        );
2391        assert_eq!(
2392            String::from_utf8(fs.read_file_sync(path!("/tmp/foo")).unwrap()).unwrap(),
2393            "zero\none\ntwo\nthree\nfour\nfive\n"
2394        );
2395        request.await.unwrap();
2396    }
2397
2398    #[gpui::test]
2399    async fn test_reading_from_line(cx: &mut TestAppContext) {
2400        init_test(cx);
2401
2402        let fs = FakeFs::new(cx.executor());
2403        fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\nfour\n"}))
2404            .await;
2405        let project = Project::test(fs.clone(), [], cx).await;
2406        project
2407            .update(cx, |project, cx| {
2408                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
2409            })
2410            .await
2411            .unwrap();
2412
2413        let connection = Rc::new(FakeAgentConnection::new());
2414
2415        let thread = cx
2416            .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2417            .await
2418            .unwrap();
2419
2420        // Whole file
2421        let content = thread
2422            .update(cx, |thread, cx| {
2423                thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
2424            })
2425            .await
2426            .unwrap();
2427
2428        assert_eq!(content, "one\ntwo\nthree\nfour\n");
2429
2430        // Only start line
2431        let content = thread
2432            .update(cx, |thread, cx| {
2433                thread.read_text_file(path!("/tmp/foo").into(), Some(3), None, false, cx)
2434            })
2435            .await
2436            .unwrap();
2437
2438        assert_eq!(content, "three\nfour\n");
2439
2440        // Only limit
2441        let content = thread
2442            .update(cx, |thread, cx| {
2443                thread.read_text_file(path!("/tmp/foo").into(), None, Some(2), false, cx)
2444            })
2445            .await
2446            .unwrap();
2447
2448        assert_eq!(content, "one\ntwo\n");
2449
2450        // Range
2451        let content = thread
2452            .update(cx, |thread, cx| {
2453                thread.read_text_file(path!("/tmp/foo").into(), Some(2), Some(2), false, cx)
2454            })
2455            .await
2456            .unwrap();
2457
2458        assert_eq!(content, "two\nthree\n");
2459
2460        // Invalid
2461        let err = thread
2462            .update(cx, |thread, cx| {
2463                thread.read_text_file(path!("/tmp/foo").into(), Some(6), Some(2), false, cx)
2464            })
2465            .await
2466            .unwrap_err();
2467
2468        assert_eq!(
2469            err.to_string(),
2470            "Invalid params: \"Attempting to read beyond the end of the file, line 5:0\""
2471        );
2472    }
2473
2474    #[gpui::test]
2475    async fn test_reading_empty_file(cx: &mut TestAppContext) {
2476        init_test(cx);
2477
2478        let fs = FakeFs::new(cx.executor());
2479        fs.insert_tree(path!("/tmp"), json!({"foo": ""})).await;
2480        let project = Project::test(fs.clone(), [], cx).await;
2481        project
2482            .update(cx, |project, cx| {
2483                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
2484            })
2485            .await
2486            .unwrap();
2487
2488        let connection = Rc::new(FakeAgentConnection::new());
2489
2490        let thread = cx
2491            .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2492            .await
2493            .unwrap();
2494
2495        // Whole file
2496        let content = thread
2497            .update(cx, |thread, cx| {
2498                thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
2499            })
2500            .await
2501            .unwrap();
2502
2503        assert_eq!(content, "");
2504
2505        // Only start line
2506        let content = thread
2507            .update(cx, |thread, cx| {
2508                thread.read_text_file(path!("/tmp/foo").into(), Some(1), None, false, cx)
2509            })
2510            .await
2511            .unwrap();
2512
2513        assert_eq!(content, "");
2514
2515        // Only limit
2516        let content = thread
2517            .update(cx, |thread, cx| {
2518                thread.read_text_file(path!("/tmp/foo").into(), None, Some(2), false, cx)
2519            })
2520            .await
2521            .unwrap();
2522
2523        assert_eq!(content, "");
2524
2525        // Range
2526        let content = thread
2527            .update(cx, |thread, cx| {
2528                thread.read_text_file(path!("/tmp/foo").into(), Some(1), Some(1), false, cx)
2529            })
2530            .await
2531            .unwrap();
2532
2533        assert_eq!(content, "");
2534
2535        // Invalid
2536        let err = thread
2537            .update(cx, |thread, cx| {
2538                thread.read_text_file(path!("/tmp/foo").into(), Some(5), Some(2), false, cx)
2539            })
2540            .await
2541            .unwrap_err();
2542
2543        assert_eq!(
2544            err.to_string(),
2545            "Invalid params: \"Attempting to read beyond the end of the file, line 1:0\""
2546        );
2547    }
2548    #[gpui::test]
2549    async fn test_reading_non_existing_file(cx: &mut TestAppContext) {
2550        init_test(cx);
2551
2552        let fs = FakeFs::new(cx.executor());
2553        fs.insert_tree(path!("/tmp"), json!({})).await;
2554        let project = Project::test(fs.clone(), [], cx).await;
2555        project
2556            .update(cx, |project, cx| {
2557                project.find_or_create_worktree(path!("/tmp"), true, cx)
2558            })
2559            .await
2560            .unwrap();
2561
2562        let connection = Rc::new(FakeAgentConnection::new());
2563
2564        let thread = cx
2565            .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2566            .await
2567            .unwrap();
2568
2569        // Out of project file
2570        let err = thread
2571            .update(cx, |thread, cx| {
2572                thread.read_text_file(path!("/foo").into(), None, None, false, cx)
2573            })
2574            .await
2575            .unwrap_err();
2576
2577        assert_eq!(err.code, acp::ErrorCode::RESOURCE_NOT_FOUND.code);
2578    }
2579
2580    #[gpui::test]
2581    async fn test_succeeding_canceled_toolcall(cx: &mut TestAppContext) {
2582        init_test(cx);
2583
2584        let fs = FakeFs::new(cx.executor());
2585        let project = Project::test(fs, [], cx).await;
2586        let id = acp::ToolCallId("test".into());
2587
2588        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2589            let id = id.clone();
2590            move |_, thread, mut cx| {
2591                let id = id.clone();
2592                async move {
2593                    thread
2594                        .update(&mut cx, |thread, cx| {
2595                            thread.handle_session_update(
2596                                acp::SessionUpdate::ToolCall(acp::ToolCall {
2597                                    id: id.clone(),
2598                                    title: "Label".into(),
2599                                    kind: acp::ToolKind::Fetch,
2600                                    status: acp::ToolCallStatus::InProgress,
2601                                    content: vec![],
2602                                    locations: vec![],
2603                                    raw_input: None,
2604                                    raw_output: None,
2605                                    meta: None,
2606                                }),
2607                                cx,
2608                            )
2609                        })
2610                        .unwrap()
2611                        .unwrap();
2612                    Ok(acp::PromptResponse {
2613                        stop_reason: acp::StopReason::EndTurn,
2614                        meta: None,
2615                    })
2616                }
2617                .boxed_local()
2618            }
2619        }));
2620
2621        let thread = cx
2622            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2623            .await
2624            .unwrap();
2625
2626        let request = thread.update(cx, |thread, cx| {
2627            thread.send_raw("Fetch https://example.com", cx)
2628        });
2629
2630        run_until_first_tool_call(&thread, cx).await;
2631
2632        thread.read_with(cx, |thread, _| {
2633            assert!(matches!(
2634                thread.entries[1],
2635                AgentThreadEntry::ToolCall(ToolCall {
2636                    status: ToolCallStatus::InProgress,
2637                    ..
2638                })
2639            ));
2640        });
2641
2642        thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2643
2644        thread.read_with(cx, |thread, _| {
2645            assert!(matches!(
2646                &thread.entries[1],
2647                AgentThreadEntry::ToolCall(ToolCall {
2648                    status: ToolCallStatus::Canceled,
2649                    ..
2650                })
2651            ));
2652        });
2653
2654        thread
2655            .update(cx, |thread, cx| {
2656                thread.handle_session_update(
2657                    acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate {
2658                        id,
2659                        fields: acp::ToolCallUpdateFields {
2660                            status: Some(acp::ToolCallStatus::Completed),
2661                            ..Default::default()
2662                        },
2663                        meta: None,
2664                    }),
2665                    cx,
2666                )
2667            })
2668            .unwrap();
2669
2670        request.await.unwrap();
2671
2672        thread.read_with(cx, |thread, _| {
2673            assert!(matches!(
2674                thread.entries[1],
2675                AgentThreadEntry::ToolCall(ToolCall {
2676                    status: ToolCallStatus::Completed,
2677                    ..
2678                })
2679            ));
2680        });
2681    }
2682
2683    #[gpui::test]
2684    async fn test_no_pending_edits_if_tool_calls_are_completed(cx: &mut TestAppContext) {
2685        init_test(cx);
2686        let fs = FakeFs::new(cx.background_executor.clone());
2687        fs.insert_tree(path!("/test"), json!({})).await;
2688        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
2689
2690        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2691            move |_, thread, mut cx| {
2692                async move {
2693                    thread
2694                        .update(&mut cx, |thread, cx| {
2695                            thread.handle_session_update(
2696                                acp::SessionUpdate::ToolCall(acp::ToolCall {
2697                                    id: acp::ToolCallId("test".into()),
2698                                    title: "Label".into(),
2699                                    kind: acp::ToolKind::Edit,
2700                                    status: acp::ToolCallStatus::Completed,
2701                                    content: vec![acp::ToolCallContent::Diff {
2702                                        diff: acp::Diff {
2703                                            path: "/test/test.txt".into(),
2704                                            old_text: None,
2705                                            new_text: "foo".into(),
2706                                            meta: None,
2707                                        },
2708                                    }],
2709                                    locations: vec![],
2710                                    raw_input: None,
2711                                    raw_output: None,
2712                                    meta: None,
2713                                }),
2714                                cx,
2715                            )
2716                        })
2717                        .unwrap()
2718                        .unwrap();
2719                    Ok(acp::PromptResponse {
2720                        stop_reason: acp::StopReason::EndTurn,
2721                        meta: None,
2722                    })
2723                }
2724                .boxed_local()
2725            }
2726        }));
2727
2728        let thread = cx
2729            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2730            .await
2731            .unwrap();
2732
2733        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Hi".into()], cx)))
2734            .await
2735            .unwrap();
2736
2737        assert!(cx.read(|cx| !thread.read(cx).has_pending_edit_tool_calls()));
2738    }
2739
2740    #[gpui::test(iterations = 10)]
2741    async fn test_checkpoints(cx: &mut TestAppContext) {
2742        init_test(cx);
2743        let fs = FakeFs::new(cx.background_executor.clone());
2744        fs.insert_tree(
2745            path!("/test"),
2746            json!({
2747                ".git": {}
2748            }),
2749        )
2750        .await;
2751        let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
2752
2753        let simulate_changes = Arc::new(AtomicBool::new(true));
2754        let next_filename = Arc::new(AtomicUsize::new(0));
2755        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2756            let simulate_changes = simulate_changes.clone();
2757            let next_filename = next_filename.clone();
2758            let fs = fs.clone();
2759            move |request, thread, mut cx| {
2760                let fs = fs.clone();
2761                let simulate_changes = simulate_changes.clone();
2762                let next_filename = next_filename.clone();
2763                async move {
2764                    if simulate_changes.load(SeqCst) {
2765                        let filename = format!("/test/file-{}", next_filename.fetch_add(1, SeqCst));
2766                        fs.write(Path::new(&filename), b"").await?;
2767                    }
2768
2769                    let acp::ContentBlock::Text(content) = &request.prompt[0] else {
2770                        panic!("expected text content block");
2771                    };
2772                    thread.update(&mut cx, |thread, cx| {
2773                        thread
2774                            .handle_session_update(
2775                                acp::SessionUpdate::AgentMessageChunk {
2776                                    content: content.text.to_uppercase().into(),
2777                                },
2778                                cx,
2779                            )
2780                            .unwrap();
2781                    })?;
2782                    Ok(acp::PromptResponse {
2783                        stop_reason: acp::StopReason::EndTurn,
2784                        meta: None,
2785                    })
2786                }
2787                .boxed_local()
2788            }
2789        }));
2790        let thread = cx
2791            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2792            .await
2793            .unwrap();
2794
2795        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Lorem".into()], cx)))
2796            .await
2797            .unwrap();
2798        thread.read_with(cx, |thread, cx| {
2799            assert_eq!(
2800                thread.to_markdown(cx),
2801                indoc! {"
2802                    ## User (checkpoint)
2803
2804                    Lorem
2805
2806                    ## Assistant
2807
2808                    LOREM
2809
2810                "}
2811            );
2812        });
2813        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2814
2815        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["ipsum".into()], cx)))
2816            .await
2817            .unwrap();
2818        thread.read_with(cx, |thread, cx| {
2819            assert_eq!(
2820                thread.to_markdown(cx),
2821                indoc! {"
2822                    ## User (checkpoint)
2823
2824                    Lorem
2825
2826                    ## Assistant
2827
2828                    LOREM
2829
2830                    ## User (checkpoint)
2831
2832                    ipsum
2833
2834                    ## Assistant
2835
2836                    IPSUM
2837
2838                "}
2839            );
2840        });
2841        assert_eq!(
2842            fs.files(),
2843            vec![
2844                Path::new(path!("/test/file-0")),
2845                Path::new(path!("/test/file-1"))
2846            ]
2847        );
2848
2849        // Checkpoint isn't stored when there are no changes.
2850        simulate_changes.store(false, SeqCst);
2851        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["dolor".into()], cx)))
2852            .await
2853            .unwrap();
2854        thread.read_with(cx, |thread, cx| {
2855            assert_eq!(
2856                thread.to_markdown(cx),
2857                indoc! {"
2858                    ## User (checkpoint)
2859
2860                    Lorem
2861
2862                    ## Assistant
2863
2864                    LOREM
2865
2866                    ## User (checkpoint)
2867
2868                    ipsum
2869
2870                    ## Assistant
2871
2872                    IPSUM
2873
2874                    ## User
2875
2876                    dolor
2877
2878                    ## Assistant
2879
2880                    DOLOR
2881
2882                "}
2883            );
2884        });
2885        assert_eq!(
2886            fs.files(),
2887            vec![
2888                Path::new(path!("/test/file-0")),
2889                Path::new(path!("/test/file-1"))
2890            ]
2891        );
2892
2893        // Rewinding the conversation truncates the history and restores the checkpoint.
2894        thread
2895            .update(cx, |thread, cx| {
2896                let AgentThreadEntry::UserMessage(message) = &thread.entries[2] else {
2897                    panic!("unexpected entries {:?}", thread.entries)
2898                };
2899                thread.restore_checkpoint(message.id.clone().unwrap(), cx)
2900            })
2901            .await
2902            .unwrap();
2903        thread.read_with(cx, |thread, cx| {
2904            assert_eq!(
2905                thread.to_markdown(cx),
2906                indoc! {"
2907                    ## User (checkpoint)
2908
2909                    Lorem
2910
2911                    ## Assistant
2912
2913                    LOREM
2914
2915                "}
2916            );
2917        });
2918        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2919    }
2920
2921    #[gpui::test]
2922    async fn test_tool_result_refusal(cx: &mut TestAppContext) {
2923        use std::sync::atomic::AtomicUsize;
2924        init_test(cx);
2925
2926        let fs = FakeFs::new(cx.executor());
2927        let project = Project::test(fs, None, cx).await;
2928
2929        // Create a connection that simulates refusal after tool result
2930        let prompt_count = Arc::new(AtomicUsize::new(0));
2931        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2932            let prompt_count = prompt_count.clone();
2933            move |_request, thread, mut cx| {
2934                let count = prompt_count.fetch_add(1, SeqCst);
2935                async move {
2936                    if count == 0 {
2937                        // First prompt: Generate a tool call with result
2938                        thread.update(&mut cx, |thread, cx| {
2939                            thread
2940                                .handle_session_update(
2941                                    acp::SessionUpdate::ToolCall(acp::ToolCall {
2942                                        id: acp::ToolCallId("tool1".into()),
2943                                        title: "Test Tool".into(),
2944                                        kind: acp::ToolKind::Fetch,
2945                                        status: acp::ToolCallStatus::Completed,
2946                                        content: vec![],
2947                                        locations: vec![],
2948                                        raw_input: Some(serde_json::json!({"query": "test"})),
2949                                        raw_output: Some(
2950                                            serde_json::json!({"result": "inappropriate content"}),
2951                                        ),
2952                                        meta: None,
2953                                    }),
2954                                    cx,
2955                                )
2956                                .unwrap();
2957                        })?;
2958
2959                        // Now return refusal because of the tool result
2960                        Ok(acp::PromptResponse {
2961                            stop_reason: acp::StopReason::Refusal,
2962                            meta: None,
2963                        })
2964                    } else {
2965                        Ok(acp::PromptResponse {
2966                            stop_reason: acp::StopReason::EndTurn,
2967                            meta: None,
2968                        })
2969                    }
2970                }
2971                .boxed_local()
2972            }
2973        }));
2974
2975        let thread = cx
2976            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2977            .await
2978            .unwrap();
2979
2980        // Track if we see a Refusal event
2981        let saw_refusal_event = Arc::new(std::sync::Mutex::new(false));
2982        let saw_refusal_event_captured = saw_refusal_event.clone();
2983        thread.update(cx, |_thread, cx| {
2984            cx.subscribe(
2985                &thread,
2986                move |_thread, _event_thread, event: &AcpThreadEvent, _cx| {
2987                    if matches!(event, AcpThreadEvent::Refusal) {
2988                        *saw_refusal_event_captured.lock().unwrap() = true;
2989                    }
2990                },
2991            )
2992            .detach();
2993        });
2994
2995        // Send a user message - this will trigger tool call and then refusal
2996        let send_task = thread.update(cx, |thread, cx| {
2997            thread.send(
2998                vec![acp::ContentBlock::Text(acp::TextContent {
2999                    text: "Hello".into(),
3000                    annotations: None,
3001                    meta: None,
3002                })],
3003                cx,
3004            )
3005        });
3006        cx.background_executor.spawn(send_task).detach();
3007        cx.run_until_parked();
3008
3009        // Verify that:
3010        // 1. A Refusal event WAS emitted (because it's a tool result refusal, not user prompt)
3011        // 2. The user message was NOT truncated
3012        assert!(
3013            *saw_refusal_event.lock().unwrap(),
3014            "Refusal event should be emitted for tool result refusals"
3015        );
3016
3017        thread.read_with(cx, |thread, _| {
3018            let entries = thread.entries();
3019            assert!(entries.len() >= 2, "Should have user message and tool call");
3020
3021            // Verify user message is still there
3022            assert!(
3023                matches!(entries[0], AgentThreadEntry::UserMessage(_)),
3024                "User message should not be truncated"
3025            );
3026
3027            // Verify tool call is there with result
3028            if let AgentThreadEntry::ToolCall(tool_call) = &entries[1] {
3029                assert!(
3030                    tool_call.raw_output.is_some(),
3031                    "Tool call should have output"
3032                );
3033            } else {
3034                panic!("Expected tool call at index 1");
3035            }
3036        });
3037    }
3038
3039    #[gpui::test]
3040    async fn test_user_prompt_refusal_emits_event(cx: &mut TestAppContext) {
3041        init_test(cx);
3042
3043        let fs = FakeFs::new(cx.executor());
3044        let project = Project::test(fs, None, cx).await;
3045
3046        let refuse_next = Arc::new(AtomicBool::new(false));
3047        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3048            let refuse_next = refuse_next.clone();
3049            move |_request, _thread, _cx| {
3050                if refuse_next.load(SeqCst) {
3051                    async move {
3052                        Ok(acp::PromptResponse {
3053                            stop_reason: acp::StopReason::Refusal,
3054                            meta: None,
3055                        })
3056                    }
3057                    .boxed_local()
3058                } else {
3059                    async move {
3060                        Ok(acp::PromptResponse {
3061                            stop_reason: acp::StopReason::EndTurn,
3062                            meta: None,
3063                        })
3064                    }
3065                    .boxed_local()
3066                }
3067            }
3068        }));
3069
3070        let thread = cx
3071            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
3072            .await
3073            .unwrap();
3074
3075        // Track if we see a Refusal event
3076        let saw_refusal_event = Arc::new(std::sync::Mutex::new(false));
3077        let saw_refusal_event_captured = saw_refusal_event.clone();
3078        thread.update(cx, |_thread, cx| {
3079            cx.subscribe(
3080                &thread,
3081                move |_thread, _event_thread, event: &AcpThreadEvent, _cx| {
3082                    if matches!(event, AcpThreadEvent::Refusal) {
3083                        *saw_refusal_event_captured.lock().unwrap() = true;
3084                    }
3085                },
3086            )
3087            .detach();
3088        });
3089
3090        // Send a message that will be refused
3091        refuse_next.store(true, SeqCst);
3092        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx)))
3093            .await
3094            .unwrap();
3095
3096        // Verify that a Refusal event WAS emitted for user prompt refusal
3097        assert!(
3098            *saw_refusal_event.lock().unwrap(),
3099            "Refusal event should be emitted for user prompt refusals"
3100        );
3101
3102        // Verify the message was truncated (user prompt refusal)
3103        thread.read_with(cx, |thread, cx| {
3104            assert_eq!(thread.to_markdown(cx), "");
3105        });
3106    }
3107
3108    #[gpui::test]
3109    async fn test_refusal(cx: &mut TestAppContext) {
3110        init_test(cx);
3111        let fs = FakeFs::new(cx.background_executor.clone());
3112        fs.insert_tree(path!("/"), json!({})).await;
3113        let project = Project::test(fs.clone(), [path!("/").as_ref()], cx).await;
3114
3115        let refuse_next = Arc::new(AtomicBool::new(false));
3116        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3117            let refuse_next = refuse_next.clone();
3118            move |request, thread, mut cx| {
3119                let refuse_next = refuse_next.clone();
3120                async move {
3121                    if refuse_next.load(SeqCst) {
3122                        return Ok(acp::PromptResponse {
3123                            stop_reason: acp::StopReason::Refusal,
3124                            meta: None,
3125                        });
3126                    }
3127
3128                    let acp::ContentBlock::Text(content) = &request.prompt[0] else {
3129                        panic!("expected text content block");
3130                    };
3131                    thread.update(&mut cx, |thread, cx| {
3132                        thread
3133                            .handle_session_update(
3134                                acp::SessionUpdate::AgentMessageChunk {
3135                                    content: content.text.to_uppercase().into(),
3136                                },
3137                                cx,
3138                            )
3139                            .unwrap();
3140                    })?;
3141                    Ok(acp::PromptResponse {
3142                        stop_reason: acp::StopReason::EndTurn,
3143                        meta: None,
3144                    })
3145                }
3146                .boxed_local()
3147            }
3148        }));
3149        let thread = cx
3150            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
3151            .await
3152            .unwrap();
3153
3154        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx)))
3155            .await
3156            .unwrap();
3157        thread.read_with(cx, |thread, cx| {
3158            assert_eq!(
3159                thread.to_markdown(cx),
3160                indoc! {"
3161                    ## User
3162
3163                    hello
3164
3165                    ## Assistant
3166
3167                    HELLO
3168
3169                "}
3170            );
3171        });
3172
3173        // Simulate refusing the second message. The message should be truncated
3174        // when a user prompt is refused.
3175        refuse_next.store(true, SeqCst);
3176        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["world".into()], cx)))
3177            .await
3178            .unwrap();
3179        thread.read_with(cx, |thread, cx| {
3180            assert_eq!(
3181                thread.to_markdown(cx),
3182                indoc! {"
3183                    ## User
3184
3185                    hello
3186
3187                    ## Assistant
3188
3189                    HELLO
3190
3191                "}
3192            );
3193        });
3194    }
3195
3196    async fn run_until_first_tool_call(
3197        thread: &Entity<AcpThread>,
3198        cx: &mut TestAppContext,
3199    ) -> usize {
3200        let (mut tx, mut rx) = mpsc::channel::<usize>(1);
3201
3202        let subscription = cx.update(|cx| {
3203            cx.subscribe(thread, move |thread, _, cx| {
3204                for (ix, entry) in thread.read(cx).entries.iter().enumerate() {
3205                    if matches!(entry, AgentThreadEntry::ToolCall(_)) {
3206                        return tx.try_send(ix).unwrap();
3207                    }
3208                }
3209            })
3210        });
3211
3212        select! {
3213            _ = futures::FutureExt::fuse(smol::Timer::after(Duration::from_secs(10))) => {
3214                panic!("Timeout waiting for tool call")
3215            }
3216            ix = rx.next().fuse() => {
3217                drop(subscription);
3218                ix.unwrap()
3219            }
3220        }
3221    }
3222
3223    #[derive(Clone, Default)]
3224    struct FakeAgentConnection {
3225        auth_methods: Vec<acp::AuthMethod>,
3226        sessions: Arc<parking_lot::Mutex<HashMap<acp::SessionId, WeakEntity<AcpThread>>>>,
3227        on_user_message: Option<
3228            Rc<
3229                dyn Fn(
3230                        acp::PromptRequest,
3231                        WeakEntity<AcpThread>,
3232                        AsyncApp,
3233                    ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
3234                    + 'static,
3235            >,
3236        >,
3237    }
3238
3239    impl FakeAgentConnection {
3240        fn new() -> Self {
3241            Self {
3242                auth_methods: Vec::new(),
3243                on_user_message: None,
3244                sessions: Arc::default(),
3245            }
3246        }
3247
3248        #[expect(unused)]
3249        fn with_auth_methods(mut self, auth_methods: Vec<acp::AuthMethod>) -> Self {
3250            self.auth_methods = auth_methods;
3251            self
3252        }
3253
3254        fn on_user_message(
3255            mut self,
3256            handler: impl Fn(
3257                acp::PromptRequest,
3258                WeakEntity<AcpThread>,
3259                AsyncApp,
3260            ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
3261            + 'static,
3262        ) -> Self {
3263            self.on_user_message.replace(Rc::new(handler));
3264            self
3265        }
3266    }
3267
3268    impl AgentConnection for FakeAgentConnection {
3269        fn auth_methods(&self) -> &[acp::AuthMethod] {
3270            &self.auth_methods
3271        }
3272
3273        fn new_thread(
3274            self: Rc<Self>,
3275            project: Entity<Project>,
3276            _cwd: &Path,
3277            cx: &mut App,
3278        ) -> Task<gpui::Result<Entity<AcpThread>>> {
3279            let session_id = acp::SessionId(
3280                rand::rng()
3281                    .sample_iter(&distr::Alphanumeric)
3282                    .take(7)
3283                    .map(char::from)
3284                    .collect::<String>()
3285                    .into(),
3286            );
3287            let action_log = cx.new(|_| ActionLog::new(project.clone()));
3288            let thread = cx.new(|cx| {
3289                AcpThread::new(
3290                    "Test",
3291                    self.clone(),
3292                    project,
3293                    action_log,
3294                    session_id.clone(),
3295                    watch::Receiver::constant(acp::PromptCapabilities {
3296                        image: true,
3297                        audio: true,
3298                        embedded_context: true,
3299                        meta: None,
3300                    }),
3301                    cx,
3302                )
3303            });
3304            self.sessions.lock().insert(session_id, thread.downgrade());
3305            Task::ready(Ok(thread))
3306        }
3307
3308        fn authenticate(&self, method: acp::AuthMethodId, _cx: &mut App) -> Task<gpui::Result<()>> {
3309            if self.auth_methods().iter().any(|m| m.id == method) {
3310                Task::ready(Ok(()))
3311            } else {
3312                Task::ready(Err(anyhow!("Invalid Auth Method")))
3313            }
3314        }
3315
3316        fn prompt(
3317            &self,
3318            _id: Option<UserMessageId>,
3319            params: acp::PromptRequest,
3320            cx: &mut App,
3321        ) -> Task<gpui::Result<acp::PromptResponse>> {
3322            let sessions = self.sessions.lock();
3323            let thread = sessions.get(&params.session_id).unwrap();
3324            if let Some(handler) = &self.on_user_message {
3325                let handler = handler.clone();
3326                let thread = thread.clone();
3327                cx.spawn(async move |cx| handler(params, thread, cx.clone()).await)
3328            } else {
3329                Task::ready(Ok(acp::PromptResponse {
3330                    stop_reason: acp::StopReason::EndTurn,
3331                    meta: None,
3332                }))
3333            }
3334        }
3335
3336        fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
3337            let sessions = self.sessions.lock();
3338            let thread = sessions.get(session_id).unwrap().clone();
3339
3340            cx.spawn(async move |cx| {
3341                thread
3342                    .update(cx, |thread, cx| thread.cancel(cx))
3343                    .unwrap()
3344                    .await
3345            })
3346            .detach();
3347        }
3348
3349        fn truncate(
3350            &self,
3351            session_id: &acp::SessionId,
3352            _cx: &App,
3353        ) -> Option<Rc<dyn AgentSessionTruncate>> {
3354            Some(Rc::new(FakeAgentSessionEditor {
3355                _session_id: session_id.clone(),
3356            }))
3357        }
3358
3359        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
3360            self
3361        }
3362    }
3363
3364    struct FakeAgentSessionEditor {
3365        _session_id: acp::SessionId,
3366    }
3367
3368    impl AgentSessionTruncate for FakeAgentSessionEditor {
3369        fn run(&self, _message_id: UserMessageId, _cx: &mut App) -> Task<Result<()>> {
3370            Task::ready(Ok(()))
3371        }
3372    }
3373
3374    #[gpui::test]
3375    async fn test_tool_call_not_found_creates_failed_entry(cx: &mut TestAppContext) {
3376        init_test(cx);
3377
3378        let fs = FakeFs::new(cx.executor());
3379        let project = Project::test(fs, [], cx).await;
3380        let connection = Rc::new(FakeAgentConnection::new());
3381        let thread = cx
3382            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
3383            .await
3384            .unwrap();
3385
3386        // Try to update a tool call that doesn't exist
3387        let nonexistent_id = acp::ToolCallId("nonexistent-tool-call".into());
3388        thread.update(cx, |thread, cx| {
3389            let result = thread.handle_session_update(
3390                acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate {
3391                    id: nonexistent_id.clone(),
3392                    fields: acp::ToolCallUpdateFields {
3393                        status: Some(acp::ToolCallStatus::Completed),
3394                        ..Default::default()
3395                    },
3396                    meta: None,
3397                }),
3398                cx,
3399            );
3400
3401            // The update should succeed (not return an error)
3402            assert!(result.is_ok());
3403
3404            // There should now be exactly one entry in the thread
3405            assert_eq!(thread.entries.len(), 1);
3406
3407            // The entry should be a failed tool call
3408            if let AgentThreadEntry::ToolCall(tool_call) = &thread.entries[0] {
3409                assert_eq!(tool_call.id, nonexistent_id);
3410                assert!(matches!(tool_call.status, ToolCallStatus::Failed));
3411                assert_eq!(tool_call.kind, acp::ToolKind::Fetch);
3412
3413                // Check that the content contains the error message
3414                assert_eq!(tool_call.content.len(), 1);
3415                if let ToolCallContent::ContentBlock(content_block) = &tool_call.content[0] {
3416                    match content_block {
3417                        ContentBlock::Markdown { markdown } => {
3418                            let markdown_text = markdown.read(cx).source();
3419                            assert!(markdown_text.contains("Tool call not found"));
3420                        }
3421                        ContentBlock::Empty => panic!("Expected markdown content, got empty"),
3422                        ContentBlock::ResourceLink { .. } => {
3423                            panic!("Expected markdown content, got resource link")
3424                        }
3425                    }
3426                } else {
3427                    panic!("Expected ContentBlock, got: {:?}", tool_call.content[0]);
3428                }
3429            } else {
3430                panic!("Expected ToolCall entry, got: {:?}", thread.entries[0]);
3431            }
3432        });
3433    }
3434}