acp_thread.rs

   1mod connection;
   2mod diff;
   3mod mention;
   4mod terminal;
   5
   6use collections::HashSet;
   7pub use connection::*;
   8pub use diff::*;
   9use language::language_settings::FormatOnSave;
  10pub use mention::*;
  11use project::lsp_store::{FormatTrigger, LspFormatTarget};
  12use serde::{Deserialize, Serialize};
  13pub use terminal::*;
  14
  15use action_log::ActionLog;
  16use agent_client_protocol as acp;
  17use anyhow::{Context as _, Result, anyhow};
  18use editor::Bias;
  19use futures::{FutureExt, channel::oneshot, future::BoxFuture};
  20use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
  21use itertools::Itertools;
  22use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, ToPoint, text_diff};
  23use markdown::Markdown;
  24use project::{AgentLocation, Project, git_store::GitStoreCheckpoint};
  25use std::collections::HashMap;
  26use std::error::Error;
  27use std::fmt::{Formatter, Write};
  28use std::ops::Range;
  29use std::process::ExitStatus;
  30use std::rc::Rc;
  31use std::time::{Duration, Instant};
  32use std::{fmt::Display, mem, path::PathBuf, sync::Arc};
  33use ui::App;
  34use util::ResultExt;
  35
  36#[derive(Debug)]
  37pub struct UserMessage {
  38    pub id: Option<UserMessageId>,
  39    pub content: ContentBlock,
  40    pub chunks: Vec<acp::ContentBlock>,
  41    pub checkpoint: Option<Checkpoint>,
  42}
  43
  44#[derive(Debug)]
  45pub struct Checkpoint {
  46    git_checkpoint: GitStoreCheckpoint,
  47    pub show: bool,
  48}
  49
  50impl UserMessage {
  51    fn to_markdown(&self, cx: &App) -> String {
  52        let mut markdown = String::new();
  53        if self
  54            .checkpoint
  55            .as_ref()
  56            .is_some_and(|checkpoint| checkpoint.show)
  57        {
  58            writeln!(markdown, "## User (checkpoint)").unwrap();
  59        } else {
  60            writeln!(markdown, "## User").unwrap();
  61        }
  62        writeln!(markdown).unwrap();
  63        writeln!(markdown, "{}", self.content.to_markdown(cx)).unwrap();
  64        writeln!(markdown).unwrap();
  65        markdown
  66    }
  67}
  68
  69#[derive(Debug, PartialEq)]
  70pub struct AssistantMessage {
  71    pub chunks: Vec<AssistantMessageChunk>,
  72}
  73
  74impl AssistantMessage {
  75    pub fn to_markdown(&self, cx: &App) -> String {
  76        format!(
  77            "## Assistant\n\n{}\n\n",
  78            self.chunks
  79                .iter()
  80                .map(|chunk| chunk.to_markdown(cx))
  81                .join("\n\n")
  82        )
  83    }
  84}
  85
  86#[derive(Debug, PartialEq)]
  87pub enum AssistantMessageChunk {
  88    Message { block: ContentBlock },
  89    Thought { block: ContentBlock },
  90}
  91
  92impl AssistantMessageChunk {
  93    pub fn from_str(chunk: &str, language_registry: &Arc<LanguageRegistry>, cx: &mut App) -> Self {
  94        Self::Message {
  95            block: ContentBlock::new(chunk.into(), language_registry, cx),
  96        }
  97    }
  98
  99    fn to_markdown(&self, cx: &App) -> String {
 100        match self {
 101            Self::Message { block } => block.to_markdown(cx).to_string(),
 102            Self::Thought { block } => {
 103                format!("<thinking>\n{}\n</thinking>", block.to_markdown(cx))
 104            }
 105        }
 106    }
 107}
 108
 109#[derive(Debug)]
 110pub enum AgentThreadEntry {
 111    UserMessage(UserMessage),
 112    AssistantMessage(AssistantMessage),
 113    ToolCall(ToolCall),
 114}
 115
 116impl AgentThreadEntry {
 117    pub fn to_markdown(&self, cx: &App) -> String {
 118        match self {
 119            Self::UserMessage(message) => message.to_markdown(cx),
 120            Self::AssistantMessage(message) => message.to_markdown(cx),
 121            Self::ToolCall(tool_call) => tool_call.to_markdown(cx),
 122        }
 123    }
 124
 125    pub fn user_message(&self) -> Option<&UserMessage> {
 126        if let AgentThreadEntry::UserMessage(message) = self {
 127            Some(message)
 128        } else {
 129            None
 130        }
 131    }
 132
 133    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 134        if let AgentThreadEntry::ToolCall(call) = self {
 135            itertools::Either::Left(call.diffs())
 136        } else {
 137            itertools::Either::Right(std::iter::empty())
 138        }
 139    }
 140
 141    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 142        if let AgentThreadEntry::ToolCall(call) = self {
 143            itertools::Either::Left(call.terminals())
 144        } else {
 145            itertools::Either::Right(std::iter::empty())
 146        }
 147    }
 148
 149    pub fn location(&self, ix: usize) -> Option<(acp::ToolCallLocation, AgentLocation)> {
 150        if let AgentThreadEntry::ToolCall(ToolCall {
 151            locations,
 152            resolved_locations,
 153            ..
 154        }) = self
 155        {
 156            Some((
 157                locations.get(ix)?.clone(),
 158                resolved_locations.get(ix)?.clone()?,
 159            ))
 160        } else {
 161            None
 162        }
 163    }
 164}
 165
 166#[derive(Debug)]
 167pub struct ToolCall {
 168    pub id: acp::ToolCallId,
 169    pub label: Entity<Markdown>,
 170    pub kind: acp::ToolKind,
 171    pub content: Vec<ToolCallContent>,
 172    pub status: ToolCallStatus,
 173    pub locations: Vec<acp::ToolCallLocation>,
 174    pub resolved_locations: Vec<Option<AgentLocation>>,
 175    pub raw_input: Option<serde_json::Value>,
 176    pub raw_output: Option<serde_json::Value>,
 177}
 178
 179impl ToolCall {
 180    fn from_acp(
 181        tool_call: acp::ToolCall,
 182        status: ToolCallStatus,
 183        language_registry: Arc<LanguageRegistry>,
 184        cx: &mut App,
 185    ) -> Self {
 186        Self {
 187            id: tool_call.id,
 188            label: cx.new(|cx| {
 189                Markdown::new(
 190                    tool_call.title.into(),
 191                    Some(language_registry.clone()),
 192                    None,
 193                    cx,
 194                )
 195            }),
 196            kind: tool_call.kind,
 197            content: tool_call
 198                .content
 199                .into_iter()
 200                .map(|content| ToolCallContent::from_acp(content, language_registry.clone(), cx))
 201                .collect(),
 202            locations: tool_call.locations,
 203            resolved_locations: Vec::default(),
 204            status,
 205            raw_input: tool_call.raw_input,
 206            raw_output: tool_call.raw_output,
 207        }
 208    }
 209
 210    fn update_fields(
 211        &mut self,
 212        fields: acp::ToolCallUpdateFields,
 213        language_registry: Arc<LanguageRegistry>,
 214        cx: &mut App,
 215    ) {
 216        let acp::ToolCallUpdateFields {
 217            kind,
 218            status,
 219            title,
 220            content,
 221            locations,
 222            raw_input,
 223            raw_output,
 224        } = fields;
 225
 226        if let Some(kind) = kind {
 227            self.kind = kind;
 228        }
 229
 230        if let Some(status) = status {
 231            self.status = status.into();
 232        }
 233
 234        if let Some(title) = title {
 235            self.label.update(cx, |label, cx| {
 236                label.replace(title, cx);
 237            });
 238        }
 239
 240        if let Some(content) = content {
 241            self.content = content
 242                .into_iter()
 243                .map(|chunk| ToolCallContent::from_acp(chunk, language_registry.clone(), cx))
 244                .collect();
 245        }
 246
 247        if let Some(locations) = locations {
 248            self.locations = locations;
 249        }
 250
 251        if let Some(raw_input) = raw_input {
 252            self.raw_input = Some(raw_input);
 253        }
 254
 255        if let Some(raw_output) = raw_output {
 256            if self.content.is_empty()
 257                && let Some(markdown) = markdown_for_raw_output(&raw_output, &language_registry, cx)
 258            {
 259                self.content
 260                    .push(ToolCallContent::ContentBlock(ContentBlock::Markdown {
 261                        markdown,
 262                    }));
 263            }
 264            self.raw_output = Some(raw_output);
 265        }
 266    }
 267
 268    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 269        self.content.iter().filter_map(|content| match content {
 270            ToolCallContent::Diff(diff) => Some(diff),
 271            ToolCallContent::ContentBlock(_) => None,
 272            ToolCallContent::Terminal(_) => None,
 273        })
 274    }
 275
 276    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 277        self.content.iter().filter_map(|content| match content {
 278            ToolCallContent::Terminal(terminal) => Some(terminal),
 279            ToolCallContent::ContentBlock(_) => None,
 280            ToolCallContent::Diff(_) => None,
 281        })
 282    }
 283
 284    fn to_markdown(&self, cx: &App) -> String {
 285        let mut markdown = format!(
 286            "**Tool Call: {}**\nStatus: {}\n\n",
 287            self.label.read(cx).source(),
 288            self.status
 289        );
 290        for content in &self.content {
 291            markdown.push_str(content.to_markdown(cx).as_str());
 292            markdown.push_str("\n\n");
 293        }
 294        markdown
 295    }
 296
 297    async fn resolve_location(
 298        location: acp::ToolCallLocation,
 299        project: WeakEntity<Project>,
 300        cx: &mut AsyncApp,
 301    ) -> Option<AgentLocation> {
 302        let buffer = project
 303            .update(cx, |project, cx| {
 304                project
 305                    .project_path_for_absolute_path(&location.path, cx)
 306                    .map(|path| project.open_buffer(path, cx))
 307            })
 308            .ok()??;
 309        let buffer = buffer.await.log_err()?;
 310        let position = buffer
 311            .update(cx, |buffer, _| {
 312                if let Some(row) = location.line {
 313                    let snapshot = buffer.snapshot();
 314                    let column = snapshot.indent_size_for_line(row).len;
 315                    let point = snapshot.clip_point(Point::new(row, column), Bias::Left);
 316                    snapshot.anchor_before(point)
 317                } else {
 318                    Anchor::MIN
 319                }
 320            })
 321            .ok()?;
 322
 323        Some(AgentLocation {
 324            buffer: buffer.downgrade(),
 325            position,
 326        })
 327    }
 328
 329    fn resolve_locations(
 330        &self,
 331        project: Entity<Project>,
 332        cx: &mut App,
 333    ) -> Task<Vec<Option<AgentLocation>>> {
 334        let locations = self.locations.clone();
 335        project.update(cx, |_, cx| {
 336            cx.spawn(async move |project, cx| {
 337                let mut new_locations = Vec::new();
 338                for location in locations {
 339                    new_locations.push(Self::resolve_location(location, project.clone(), cx).await);
 340                }
 341                new_locations
 342            })
 343        })
 344    }
 345}
 346
 347#[derive(Debug)]
 348pub enum ToolCallStatus {
 349    /// The tool call hasn't started running yet, but we start showing it to
 350    /// the user.
 351    Pending,
 352    /// The tool call is waiting for confirmation from the user.
 353    WaitingForConfirmation {
 354        options: Vec<acp::PermissionOption>,
 355        respond_tx: oneshot::Sender<acp::PermissionOptionId>,
 356    },
 357    /// The tool call is currently running.
 358    InProgress,
 359    /// The tool call completed successfully.
 360    Completed,
 361    /// The tool call failed.
 362    Failed,
 363    /// The user rejected the tool call.
 364    Rejected,
 365    /// The user canceled generation so the tool call was canceled.
 366    Canceled,
 367}
 368
 369impl From<acp::ToolCallStatus> for ToolCallStatus {
 370    fn from(status: acp::ToolCallStatus) -> Self {
 371        match status {
 372            acp::ToolCallStatus::Pending => Self::Pending,
 373            acp::ToolCallStatus::InProgress => Self::InProgress,
 374            acp::ToolCallStatus::Completed => Self::Completed,
 375            acp::ToolCallStatus::Failed => Self::Failed,
 376        }
 377    }
 378}
 379
 380impl Display for ToolCallStatus {
 381    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 382        write!(
 383            f,
 384            "{}",
 385            match self {
 386                ToolCallStatus::Pending => "Pending",
 387                ToolCallStatus::WaitingForConfirmation { .. } => "Waiting for confirmation",
 388                ToolCallStatus::InProgress => "In Progress",
 389                ToolCallStatus::Completed => "Completed",
 390                ToolCallStatus::Failed => "Failed",
 391                ToolCallStatus::Rejected => "Rejected",
 392                ToolCallStatus::Canceled => "Canceled",
 393            }
 394        )
 395    }
 396}
 397
 398#[derive(Debug, PartialEq, Clone)]
 399pub enum ContentBlock {
 400    Empty,
 401    Markdown { markdown: Entity<Markdown> },
 402    ResourceLink { resource_link: acp::ResourceLink },
 403}
 404
 405impl ContentBlock {
 406    pub fn new(
 407        block: acp::ContentBlock,
 408        language_registry: &Arc<LanguageRegistry>,
 409        cx: &mut App,
 410    ) -> Self {
 411        let mut this = Self::Empty;
 412        this.append(block, language_registry, cx);
 413        this
 414    }
 415
 416    pub fn new_combined(
 417        blocks: impl IntoIterator<Item = acp::ContentBlock>,
 418        language_registry: Arc<LanguageRegistry>,
 419        cx: &mut App,
 420    ) -> Self {
 421        let mut this = Self::Empty;
 422        for block in blocks {
 423            this.append(block, &language_registry, cx);
 424        }
 425        this
 426    }
 427
 428    pub fn append(
 429        &mut self,
 430        block: acp::ContentBlock,
 431        language_registry: &Arc<LanguageRegistry>,
 432        cx: &mut App,
 433    ) {
 434        if matches!(self, ContentBlock::Empty)
 435            && let acp::ContentBlock::ResourceLink(resource_link) = block
 436        {
 437            *self = ContentBlock::ResourceLink { resource_link };
 438            return;
 439        }
 440
 441        let new_content = self.block_string_contents(block);
 442
 443        match self {
 444            ContentBlock::Empty => {
 445                *self = Self::create_markdown_block(new_content, language_registry, cx);
 446            }
 447            ContentBlock::Markdown { markdown } => {
 448                markdown.update(cx, |markdown, cx| markdown.append(&new_content, cx));
 449            }
 450            ContentBlock::ResourceLink { resource_link } => {
 451                let existing_content = Self::resource_link_md(&resource_link.uri);
 452                let combined = format!("{}\n{}", existing_content, new_content);
 453
 454                *self = Self::create_markdown_block(combined, language_registry, cx);
 455            }
 456        }
 457    }
 458
 459    fn create_markdown_block(
 460        content: String,
 461        language_registry: &Arc<LanguageRegistry>,
 462        cx: &mut App,
 463    ) -> ContentBlock {
 464        ContentBlock::Markdown {
 465            markdown: cx
 466                .new(|cx| Markdown::new(content.into(), Some(language_registry.clone()), None, cx)),
 467        }
 468    }
 469
 470    fn block_string_contents(&self, block: acp::ContentBlock) -> String {
 471        match block {
 472            acp::ContentBlock::Text(text_content) => text_content.text,
 473            acp::ContentBlock::ResourceLink(resource_link) => {
 474                Self::resource_link_md(&resource_link.uri)
 475            }
 476            acp::ContentBlock::Resource(acp::EmbeddedResource {
 477                resource:
 478                    acp::EmbeddedResourceResource::TextResourceContents(acp::TextResourceContents {
 479                        uri,
 480                        ..
 481                    }),
 482                ..
 483            }) => Self::resource_link_md(&uri),
 484            acp::ContentBlock::Image(image) => Self::image_md(&image),
 485            acp::ContentBlock::Audio(_) | acp::ContentBlock::Resource(_) => String::new(),
 486        }
 487    }
 488
 489    fn resource_link_md(uri: &str) -> String {
 490        if let Some(uri) = MentionUri::parse(uri).log_err() {
 491            uri.as_link().to_string()
 492        } else {
 493            uri.to_string()
 494        }
 495    }
 496
 497    fn image_md(_image: &acp::ImageContent) -> String {
 498        "`Image`".into()
 499    }
 500
 501    fn to_markdown<'a>(&'a self, cx: &'a App) -> &'a str {
 502        match self {
 503            ContentBlock::Empty => "",
 504            ContentBlock::Markdown { markdown } => markdown.read(cx).source(),
 505            ContentBlock::ResourceLink { resource_link } => &resource_link.uri,
 506        }
 507    }
 508
 509    pub fn markdown(&self) -> Option<&Entity<Markdown>> {
 510        match self {
 511            ContentBlock::Empty => None,
 512            ContentBlock::Markdown { markdown } => Some(markdown),
 513            ContentBlock::ResourceLink { .. } => None,
 514        }
 515    }
 516
 517    pub fn resource_link(&self) -> Option<&acp::ResourceLink> {
 518        match self {
 519            ContentBlock::ResourceLink { resource_link } => Some(resource_link),
 520            _ => None,
 521        }
 522    }
 523}
 524
 525#[derive(Debug)]
 526pub enum ToolCallContent {
 527    ContentBlock(ContentBlock),
 528    Diff(Entity<Diff>),
 529    Terminal(Entity<Terminal>),
 530}
 531
 532impl ToolCallContent {
 533    pub fn from_acp(
 534        content: acp::ToolCallContent,
 535        language_registry: Arc<LanguageRegistry>,
 536        cx: &mut App,
 537    ) -> Self {
 538        match content {
 539            acp::ToolCallContent::Content { content } => {
 540                Self::ContentBlock(ContentBlock::new(content, &language_registry, cx))
 541            }
 542            acp::ToolCallContent::Diff { diff } => Self::Diff(cx.new(|cx| {
 543                Diff::finalized(
 544                    diff.path,
 545                    diff.old_text,
 546                    diff.new_text,
 547                    language_registry,
 548                    cx,
 549                )
 550            })),
 551        }
 552    }
 553
 554    pub fn to_markdown(&self, cx: &App) -> String {
 555        match self {
 556            Self::ContentBlock(content) => content.to_markdown(cx).to_string(),
 557            Self::Diff(diff) => diff.read(cx).to_markdown(cx),
 558            Self::Terminal(terminal) => terminal.read(cx).to_markdown(cx),
 559        }
 560    }
 561}
 562
 563#[derive(Debug, PartialEq)]
 564pub enum ToolCallUpdate {
 565    UpdateFields(acp::ToolCallUpdate),
 566    UpdateDiff(ToolCallUpdateDiff),
 567    UpdateTerminal(ToolCallUpdateTerminal),
 568}
 569
 570impl ToolCallUpdate {
 571    fn id(&self) -> &acp::ToolCallId {
 572        match self {
 573            Self::UpdateFields(update) => &update.id,
 574            Self::UpdateDiff(diff) => &diff.id,
 575            Self::UpdateTerminal(terminal) => &terminal.id,
 576        }
 577    }
 578}
 579
 580impl From<acp::ToolCallUpdate> for ToolCallUpdate {
 581    fn from(update: acp::ToolCallUpdate) -> Self {
 582        Self::UpdateFields(update)
 583    }
 584}
 585
 586impl From<ToolCallUpdateDiff> for ToolCallUpdate {
 587    fn from(diff: ToolCallUpdateDiff) -> Self {
 588        Self::UpdateDiff(diff)
 589    }
 590}
 591
 592#[derive(Debug, PartialEq)]
 593pub struct ToolCallUpdateDiff {
 594    pub id: acp::ToolCallId,
 595    pub diff: Entity<Diff>,
 596}
 597
 598impl From<ToolCallUpdateTerminal> for ToolCallUpdate {
 599    fn from(terminal: ToolCallUpdateTerminal) -> Self {
 600        Self::UpdateTerminal(terminal)
 601    }
 602}
 603
 604#[derive(Debug, PartialEq)]
 605pub struct ToolCallUpdateTerminal {
 606    pub id: acp::ToolCallId,
 607    pub terminal: Entity<Terminal>,
 608}
 609
 610#[derive(Debug, Default)]
 611pub struct Plan {
 612    pub entries: Vec<PlanEntry>,
 613}
 614
 615#[derive(Debug)]
 616pub struct PlanStats<'a> {
 617    pub in_progress_entry: Option<&'a PlanEntry>,
 618    pub pending: u32,
 619    pub completed: u32,
 620}
 621
 622impl Plan {
 623    pub fn is_empty(&self) -> bool {
 624        self.entries.is_empty()
 625    }
 626
 627    pub fn stats(&self) -> PlanStats<'_> {
 628        let mut stats = PlanStats {
 629            in_progress_entry: None,
 630            pending: 0,
 631            completed: 0,
 632        };
 633
 634        for entry in &self.entries {
 635            match &entry.status {
 636                acp::PlanEntryStatus::Pending => {
 637                    stats.pending += 1;
 638                }
 639                acp::PlanEntryStatus::InProgress => {
 640                    stats.in_progress_entry = stats.in_progress_entry.or(Some(entry));
 641                }
 642                acp::PlanEntryStatus::Completed => {
 643                    stats.completed += 1;
 644                }
 645            }
 646        }
 647
 648        stats
 649    }
 650}
 651
 652#[derive(Debug)]
 653pub struct PlanEntry {
 654    pub content: Entity<Markdown>,
 655    pub priority: acp::PlanEntryPriority,
 656    pub status: acp::PlanEntryStatus,
 657}
 658
 659impl PlanEntry {
 660    pub fn from_acp(entry: acp::PlanEntry, cx: &mut App) -> Self {
 661        Self {
 662            content: cx.new(|cx| Markdown::new(entry.content.into(), None, None, cx)),
 663            priority: entry.priority,
 664            status: entry.status,
 665        }
 666    }
 667}
 668
 669#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 670pub struct TokenUsage {
 671    pub max_tokens: u64,
 672    pub used_tokens: u64,
 673}
 674
 675impl TokenUsage {
 676    pub fn ratio(&self) -> TokenUsageRatio {
 677        #[cfg(debug_assertions)]
 678        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
 679            .unwrap_or("0.8".to_string())
 680            .parse()
 681            .unwrap();
 682        #[cfg(not(debug_assertions))]
 683        let warning_threshold: f32 = 0.8;
 684
 685        // When the maximum is unknown because there is no selected model,
 686        // avoid showing the token limit warning.
 687        if self.max_tokens == 0 {
 688            TokenUsageRatio::Normal
 689        } else if self.used_tokens >= self.max_tokens {
 690            TokenUsageRatio::Exceeded
 691        } else if self.used_tokens as f32 / self.max_tokens as f32 >= warning_threshold {
 692            TokenUsageRatio::Warning
 693        } else {
 694            TokenUsageRatio::Normal
 695        }
 696    }
 697}
 698
 699#[derive(Debug, Clone, PartialEq, Eq)]
 700pub enum TokenUsageRatio {
 701    Normal,
 702    Warning,
 703    Exceeded,
 704}
 705
 706#[derive(Debug, Clone)]
 707pub struct RetryStatus {
 708    pub last_error: SharedString,
 709    pub attempt: usize,
 710    pub max_attempts: usize,
 711    pub started_at: Instant,
 712    pub duration: Duration,
 713}
 714
 715pub struct AcpThread {
 716    title: SharedString,
 717    entries: Vec<AgentThreadEntry>,
 718    plan: Plan,
 719    project: Entity<Project>,
 720    action_log: Entity<ActionLog>,
 721    shared_buffers: HashMap<Entity<Buffer>, BufferSnapshot>,
 722    send_task: Option<Task<()>>,
 723    connection: Rc<dyn AgentConnection>,
 724    session_id: acp::SessionId,
 725    token_usage: Option<TokenUsage>,
 726}
 727
 728#[derive(Debug)]
 729pub enum AcpThreadEvent {
 730    NewEntry,
 731    TitleUpdated,
 732    TokenUsageUpdated,
 733    EntryUpdated(usize),
 734    EntriesRemoved(Range<usize>),
 735    ToolAuthorizationRequired,
 736    Retry(RetryStatus),
 737    Stopped,
 738    Error,
 739    LoadError(LoadError),
 740}
 741
 742impl EventEmitter<AcpThreadEvent> for AcpThread {}
 743
 744#[derive(PartialEq, Eq)]
 745pub enum ThreadStatus {
 746    Idle,
 747    WaitingForToolConfirmation,
 748    Generating,
 749}
 750
 751#[derive(Debug, Clone)]
 752pub enum LoadError {
 753    NotInstalled {
 754        error_message: SharedString,
 755        install_message: SharedString,
 756        install_command: String,
 757    },
 758    Unsupported {
 759        error_message: SharedString,
 760        upgrade_message: SharedString,
 761        upgrade_command: String,
 762    },
 763    Exited {
 764        status: ExitStatus,
 765    },
 766    Other(SharedString),
 767}
 768
 769impl Display for LoadError {
 770    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 771        match self {
 772            LoadError::NotInstalled { error_message, .. }
 773            | LoadError::Unsupported { error_message, .. } => {
 774                write!(f, "{error_message}")
 775            }
 776            LoadError::Exited { status } => write!(f, "Server exited with status {status}"),
 777            LoadError::Other(msg) => write!(f, "{}", msg),
 778        }
 779    }
 780}
 781
 782impl Error for LoadError {}
 783
 784impl AcpThread {
 785    pub fn new(
 786        title: impl Into<SharedString>,
 787        connection: Rc<dyn AgentConnection>,
 788        project: Entity<Project>,
 789        action_log: Entity<ActionLog>,
 790        session_id: acp::SessionId,
 791    ) -> Self {
 792        Self {
 793            action_log,
 794            shared_buffers: Default::default(),
 795            entries: Default::default(),
 796            plan: Default::default(),
 797            title: title.into(),
 798            project,
 799            send_task: None,
 800            connection,
 801            session_id,
 802            token_usage: None,
 803        }
 804    }
 805
 806    pub fn connection(&self) -> &Rc<dyn AgentConnection> {
 807        &self.connection
 808    }
 809
 810    pub fn action_log(&self) -> &Entity<ActionLog> {
 811        &self.action_log
 812    }
 813
 814    pub fn project(&self) -> &Entity<Project> {
 815        &self.project
 816    }
 817
 818    pub fn title(&self) -> SharedString {
 819        self.title.clone()
 820    }
 821
 822    pub fn entries(&self) -> &[AgentThreadEntry] {
 823        &self.entries
 824    }
 825
 826    pub fn session_id(&self) -> &acp::SessionId {
 827        &self.session_id
 828    }
 829
 830    pub fn status(&self) -> ThreadStatus {
 831        if self.send_task.is_some() {
 832            if self.waiting_for_tool_confirmation() {
 833                ThreadStatus::WaitingForToolConfirmation
 834            } else {
 835                ThreadStatus::Generating
 836            }
 837        } else {
 838            ThreadStatus::Idle
 839        }
 840    }
 841
 842    pub fn token_usage(&self) -> Option<&TokenUsage> {
 843        self.token_usage.as_ref()
 844    }
 845
 846    pub fn has_pending_edit_tool_calls(&self) -> bool {
 847        for entry in self.entries.iter().rev() {
 848            match entry {
 849                AgentThreadEntry::UserMessage(_) => return false,
 850                AgentThreadEntry::ToolCall(
 851                    call @ ToolCall {
 852                        status: ToolCallStatus::InProgress | ToolCallStatus::Pending,
 853                        ..
 854                    },
 855                ) if call.diffs().next().is_some() => {
 856                    return true;
 857                }
 858                AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
 859            }
 860        }
 861
 862        false
 863    }
 864
 865    pub fn used_tools_since_last_user_message(&self) -> bool {
 866        for entry in self.entries.iter().rev() {
 867            match entry {
 868                AgentThreadEntry::UserMessage(..) => return false,
 869                AgentThreadEntry::AssistantMessage(..) => continue,
 870                AgentThreadEntry::ToolCall(..) => return true,
 871            }
 872        }
 873
 874        false
 875    }
 876
 877    pub fn handle_session_update(
 878        &mut self,
 879        update: acp::SessionUpdate,
 880        cx: &mut Context<Self>,
 881    ) -> Result<(), acp::Error> {
 882        match update {
 883            acp::SessionUpdate::UserMessageChunk { content } => {
 884                self.push_user_content_block(None, content, cx);
 885            }
 886            acp::SessionUpdate::AgentMessageChunk { content } => {
 887                self.push_assistant_content_block(content, false, cx);
 888            }
 889            acp::SessionUpdate::AgentThoughtChunk { content } => {
 890                self.push_assistant_content_block(content, true, cx);
 891            }
 892            acp::SessionUpdate::ToolCall(tool_call) => {
 893                self.upsert_tool_call(tool_call, cx)?;
 894            }
 895            acp::SessionUpdate::ToolCallUpdate(tool_call_update) => {
 896                self.update_tool_call(tool_call_update, cx)?;
 897            }
 898            acp::SessionUpdate::Plan(plan) => {
 899                self.update_plan(plan, cx);
 900            }
 901        }
 902        Ok(())
 903    }
 904
 905    pub fn push_user_content_block(
 906        &mut self,
 907        message_id: Option<UserMessageId>,
 908        chunk: acp::ContentBlock,
 909        cx: &mut Context<Self>,
 910    ) {
 911        let language_registry = self.project.read(cx).languages().clone();
 912        let entries_len = self.entries.len();
 913
 914        if let Some(last_entry) = self.entries.last_mut()
 915            && let AgentThreadEntry::UserMessage(UserMessage {
 916                id,
 917                content,
 918                chunks,
 919                ..
 920            }) = last_entry
 921        {
 922            *id = message_id.or(id.take());
 923            content.append(chunk.clone(), &language_registry, cx);
 924            chunks.push(chunk);
 925            let idx = entries_len - 1;
 926            cx.emit(AcpThreadEvent::EntryUpdated(idx));
 927        } else {
 928            let content = ContentBlock::new(chunk.clone(), &language_registry, cx);
 929            self.push_entry(
 930                AgentThreadEntry::UserMessage(UserMessage {
 931                    id: message_id,
 932                    content,
 933                    chunks: vec![chunk],
 934                    checkpoint: None,
 935                }),
 936                cx,
 937            );
 938        }
 939    }
 940
 941    pub fn push_assistant_content_block(
 942        &mut self,
 943        chunk: acp::ContentBlock,
 944        is_thought: bool,
 945        cx: &mut Context<Self>,
 946    ) {
 947        let language_registry = self.project.read(cx).languages().clone();
 948        let entries_len = self.entries.len();
 949        if let Some(last_entry) = self.entries.last_mut()
 950            && let AgentThreadEntry::AssistantMessage(AssistantMessage { chunks }) = last_entry
 951        {
 952            let idx = entries_len - 1;
 953            cx.emit(AcpThreadEvent::EntryUpdated(idx));
 954            match (chunks.last_mut(), is_thought) {
 955                (Some(AssistantMessageChunk::Message { block }), false)
 956                | (Some(AssistantMessageChunk::Thought { block }), true) => {
 957                    block.append(chunk, &language_registry, cx)
 958                }
 959                _ => {
 960                    let block = ContentBlock::new(chunk, &language_registry, cx);
 961                    if is_thought {
 962                        chunks.push(AssistantMessageChunk::Thought { block })
 963                    } else {
 964                        chunks.push(AssistantMessageChunk::Message { block })
 965                    }
 966                }
 967            }
 968        } else {
 969            let block = ContentBlock::new(chunk, &language_registry, cx);
 970            let chunk = if is_thought {
 971                AssistantMessageChunk::Thought { block }
 972            } else {
 973                AssistantMessageChunk::Message { block }
 974            };
 975
 976            self.push_entry(
 977                AgentThreadEntry::AssistantMessage(AssistantMessage {
 978                    chunks: vec![chunk],
 979                }),
 980                cx,
 981            );
 982        }
 983    }
 984
 985    fn push_entry(&mut self, entry: AgentThreadEntry, cx: &mut Context<Self>) {
 986        self.entries.push(entry);
 987        cx.emit(AcpThreadEvent::NewEntry);
 988    }
 989
 990    pub fn update_title(&mut self, title: SharedString, cx: &mut Context<Self>) -> Result<()> {
 991        self.title = title;
 992        cx.emit(AcpThreadEvent::TitleUpdated);
 993        Ok(())
 994    }
 995
 996    pub fn update_token_usage(&mut self, usage: Option<TokenUsage>, cx: &mut Context<Self>) {
 997        self.token_usage = usage;
 998        cx.emit(AcpThreadEvent::TokenUsageUpdated);
 999    }
1000
1001    pub fn update_retry_status(&mut self, status: RetryStatus, cx: &mut Context<Self>) {
1002        cx.emit(AcpThreadEvent::Retry(status));
1003    }
1004
1005    pub fn update_tool_call(
1006        &mut self,
1007        update: impl Into<ToolCallUpdate>,
1008        cx: &mut Context<Self>,
1009    ) -> Result<()> {
1010        let update = update.into();
1011        let languages = self.project.read(cx).languages().clone();
1012
1013        let (ix, current_call) = self
1014            .tool_call_mut(update.id())
1015            .context("Tool call not found")?;
1016        match update {
1017            ToolCallUpdate::UpdateFields(update) => {
1018                let location_updated = update.fields.locations.is_some();
1019                current_call.update_fields(update.fields, languages, cx);
1020                if location_updated {
1021                    self.resolve_locations(update.id, cx);
1022                }
1023            }
1024            ToolCallUpdate::UpdateDiff(update) => {
1025                current_call.content.clear();
1026                current_call
1027                    .content
1028                    .push(ToolCallContent::Diff(update.diff));
1029            }
1030            ToolCallUpdate::UpdateTerminal(update) => {
1031                current_call.content.clear();
1032                current_call
1033                    .content
1034                    .push(ToolCallContent::Terminal(update.terminal));
1035            }
1036        }
1037
1038        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1039
1040        Ok(())
1041    }
1042
1043    /// Updates a tool call if id matches an existing entry, otherwise inserts a new one.
1044    pub fn upsert_tool_call(
1045        &mut self,
1046        tool_call: acp::ToolCall,
1047        cx: &mut Context<Self>,
1048    ) -> Result<(), acp::Error> {
1049        let status = tool_call.status.into();
1050        self.upsert_tool_call_inner(tool_call.into(), status, cx)
1051    }
1052
1053    /// Fails if id does not match an existing entry.
1054    pub fn upsert_tool_call_inner(
1055        &mut self,
1056        tool_call_update: acp::ToolCallUpdate,
1057        status: ToolCallStatus,
1058        cx: &mut Context<Self>,
1059    ) -> Result<(), acp::Error> {
1060        let language_registry = self.project.read(cx).languages().clone();
1061        let id = tool_call_update.id.clone();
1062
1063        if let Some((ix, current_call)) = self.tool_call_mut(&id) {
1064            current_call.update_fields(tool_call_update.fields, language_registry, cx);
1065            current_call.status = status;
1066
1067            cx.emit(AcpThreadEvent::EntryUpdated(ix));
1068        } else {
1069            let call =
1070                ToolCall::from_acp(tool_call_update.try_into()?, status, language_registry, cx);
1071            self.push_entry(AgentThreadEntry::ToolCall(call), cx);
1072        };
1073
1074        self.resolve_locations(id, cx);
1075        Ok(())
1076    }
1077
1078    fn tool_call_mut(&mut self, id: &acp::ToolCallId) -> Option<(usize, &mut ToolCall)> {
1079        // The tool call we are looking for is typically the last one, or very close to the end.
1080        // At the moment, it doesn't seem like a hashmap would be a good fit for this use case.
1081        self.entries
1082            .iter_mut()
1083            .enumerate()
1084            .rev()
1085            .find_map(|(index, tool_call)| {
1086                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1087                    && &tool_call.id == id
1088                {
1089                    Some((index, tool_call))
1090                } else {
1091                    None
1092                }
1093            })
1094    }
1095
1096    pub fn tool_call(&mut self, id: &acp::ToolCallId) -> Option<(usize, &ToolCall)> {
1097        self.entries
1098            .iter()
1099            .enumerate()
1100            .rev()
1101            .find_map(|(index, tool_call)| {
1102                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1103                    && &tool_call.id == id
1104                {
1105                    Some((index, tool_call))
1106                } else {
1107                    None
1108                }
1109            })
1110    }
1111
1112    pub fn resolve_locations(&mut self, id: acp::ToolCallId, cx: &mut Context<Self>) {
1113        let project = self.project.clone();
1114        let Some((_, tool_call)) = self.tool_call_mut(&id) else {
1115            return;
1116        };
1117        let task = tool_call.resolve_locations(project, cx);
1118        cx.spawn(async move |this, cx| {
1119            let resolved_locations = task.await;
1120            this.update(cx, |this, cx| {
1121                let project = this.project.clone();
1122                let Some((ix, tool_call)) = this.tool_call_mut(&id) else {
1123                    return;
1124                };
1125                if let Some(Some(location)) = resolved_locations.last() {
1126                    project.update(cx, |project, cx| {
1127                        if let Some(agent_location) = project.agent_location() {
1128                            let should_ignore = agent_location.buffer == location.buffer
1129                                && location
1130                                    .buffer
1131                                    .update(cx, |buffer, _| {
1132                                        let snapshot = buffer.snapshot();
1133                                        let old_position =
1134                                            agent_location.position.to_point(&snapshot);
1135                                        let new_position = location.position.to_point(&snapshot);
1136                                        // ignore this so that when we get updates from the edit tool
1137                                        // the position doesn't reset to the startof line
1138                                        old_position.row == new_position.row
1139                                            && old_position.column > new_position.column
1140                                    })
1141                                    .ok()
1142                                    .unwrap_or_default();
1143                            if !should_ignore {
1144                                project.set_agent_location(Some(location.clone()), cx);
1145                            }
1146                        }
1147                    });
1148                }
1149                if tool_call.resolved_locations != resolved_locations {
1150                    tool_call.resolved_locations = resolved_locations;
1151                    cx.emit(AcpThreadEvent::EntryUpdated(ix));
1152                }
1153            })
1154        })
1155        .detach();
1156    }
1157
1158    pub fn request_tool_call_authorization(
1159        &mut self,
1160        tool_call: acp::ToolCallUpdate,
1161        options: Vec<acp::PermissionOption>,
1162        cx: &mut Context<Self>,
1163    ) -> Result<oneshot::Receiver<acp::PermissionOptionId>, acp::Error> {
1164        let (tx, rx) = oneshot::channel();
1165
1166        let status = ToolCallStatus::WaitingForConfirmation {
1167            options,
1168            respond_tx: tx,
1169        };
1170
1171        self.upsert_tool_call_inner(tool_call, status, cx)?;
1172        cx.emit(AcpThreadEvent::ToolAuthorizationRequired);
1173        Ok(rx)
1174    }
1175
1176    pub fn authorize_tool_call(
1177        &mut self,
1178        id: acp::ToolCallId,
1179        option_id: acp::PermissionOptionId,
1180        option_kind: acp::PermissionOptionKind,
1181        cx: &mut Context<Self>,
1182    ) {
1183        let Some((ix, call)) = self.tool_call_mut(&id) else {
1184            return;
1185        };
1186
1187        let new_status = match option_kind {
1188            acp::PermissionOptionKind::RejectOnce | acp::PermissionOptionKind::RejectAlways => {
1189                ToolCallStatus::Rejected
1190            }
1191            acp::PermissionOptionKind::AllowOnce | acp::PermissionOptionKind::AllowAlways => {
1192                ToolCallStatus::InProgress
1193            }
1194        };
1195
1196        let curr_status = mem::replace(&mut call.status, new_status);
1197
1198        if let ToolCallStatus::WaitingForConfirmation { respond_tx, .. } = curr_status {
1199            respond_tx.send(option_id).log_err();
1200        } else if cfg!(debug_assertions) {
1201            panic!("tried to authorize an already authorized tool call");
1202        }
1203
1204        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1205    }
1206
1207    /// Returns true if the last turn is awaiting tool authorization
1208    pub fn waiting_for_tool_confirmation(&self) -> bool {
1209        for entry in self.entries.iter().rev() {
1210            match &entry {
1211                AgentThreadEntry::ToolCall(call) => match call.status {
1212                    ToolCallStatus::WaitingForConfirmation { .. } => return true,
1213                    ToolCallStatus::Pending
1214                    | ToolCallStatus::InProgress
1215                    | ToolCallStatus::Completed
1216                    | ToolCallStatus::Failed
1217                    | ToolCallStatus::Rejected
1218                    | ToolCallStatus::Canceled => continue,
1219                },
1220                AgentThreadEntry::UserMessage(_) | AgentThreadEntry::AssistantMessage(_) => {
1221                    // Reached the beginning of the turn
1222                    return false;
1223                }
1224            }
1225        }
1226        false
1227    }
1228
1229    pub fn plan(&self) -> &Plan {
1230        &self.plan
1231    }
1232
1233    pub fn update_plan(&mut self, request: acp::Plan, cx: &mut Context<Self>) {
1234        let new_entries_len = request.entries.len();
1235        let mut new_entries = request.entries.into_iter();
1236
1237        // Reuse existing markdown to prevent flickering
1238        for (old, new) in self.plan.entries.iter_mut().zip(new_entries.by_ref()) {
1239            let PlanEntry {
1240                content,
1241                priority,
1242                status,
1243            } = old;
1244            content.update(cx, |old, cx| {
1245                old.replace(new.content, cx);
1246            });
1247            *priority = new.priority;
1248            *status = new.status;
1249        }
1250        for new in new_entries {
1251            self.plan.entries.push(PlanEntry::from_acp(new, cx))
1252        }
1253        self.plan.entries.truncate(new_entries_len);
1254
1255        cx.notify();
1256    }
1257
1258    fn clear_completed_plan_entries(&mut self, cx: &mut Context<Self>) {
1259        self.plan
1260            .entries
1261            .retain(|entry| !matches!(entry.status, acp::PlanEntryStatus::Completed));
1262        cx.notify();
1263    }
1264
1265    #[cfg(any(test, feature = "test-support"))]
1266    pub fn send_raw(
1267        &mut self,
1268        message: &str,
1269        cx: &mut Context<Self>,
1270    ) -> BoxFuture<'static, Result<()>> {
1271        self.send(
1272            vec![acp::ContentBlock::Text(acp::TextContent {
1273                text: message.to_string(),
1274                annotations: None,
1275            })],
1276            cx,
1277        )
1278    }
1279
1280    pub fn send(
1281        &mut self,
1282        message: Vec<acp::ContentBlock>,
1283        cx: &mut Context<Self>,
1284    ) -> BoxFuture<'static, Result<()>> {
1285        let block = ContentBlock::new_combined(
1286            message.clone(),
1287            self.project.read(cx).languages().clone(),
1288            cx,
1289        );
1290        let request = acp::PromptRequest {
1291            prompt: message.clone(),
1292            session_id: self.session_id.clone(),
1293        };
1294        let git_store = self.project.read(cx).git_store().clone();
1295
1296        let message_id = if self
1297            .connection
1298            .session_editor(&self.session_id, cx)
1299            .is_some()
1300        {
1301            Some(UserMessageId::new())
1302        } else {
1303            None
1304        };
1305
1306        self.run_turn(cx, async move |this, cx| {
1307            this.update(cx, |this, cx| {
1308                this.push_entry(
1309                    AgentThreadEntry::UserMessage(UserMessage {
1310                        id: message_id.clone(),
1311                        content: block,
1312                        chunks: message,
1313                        checkpoint: None,
1314                    }),
1315                    cx,
1316                );
1317            })
1318            .ok();
1319
1320            let old_checkpoint = git_store
1321                .update(cx, |git, cx| git.checkpoint(cx))?
1322                .await
1323                .context("failed to get old checkpoint")
1324                .log_err();
1325            this.update(cx, |this, cx| {
1326                if let Some((_ix, message)) = this.last_user_message() {
1327                    message.checkpoint = old_checkpoint.map(|git_checkpoint| Checkpoint {
1328                        git_checkpoint,
1329                        show: false,
1330                    });
1331                }
1332                this.connection.prompt(message_id, request, cx)
1333            })?
1334            .await
1335        })
1336    }
1337
1338    pub fn resume(&mut self, cx: &mut Context<Self>) -> BoxFuture<'static, Result<()>> {
1339        self.run_turn(cx, async move |this, cx| {
1340            this.update(cx, |this, cx| {
1341                this.connection
1342                    .resume(&this.session_id, cx)
1343                    .map(|resume| resume.run(cx))
1344            })?
1345            .context("resuming a session is not supported")?
1346            .await
1347        })
1348    }
1349
1350    fn run_turn(
1351        &mut self,
1352        cx: &mut Context<Self>,
1353        f: impl 'static + AsyncFnOnce(WeakEntity<Self>, &mut AsyncApp) -> Result<acp::PromptResponse>,
1354    ) -> BoxFuture<'static, Result<()>> {
1355        self.clear_completed_plan_entries(cx);
1356
1357        let (tx, rx) = oneshot::channel();
1358        let cancel_task = self.cancel(cx);
1359
1360        self.send_task = Some(cx.spawn(async move |this, cx| {
1361            cancel_task.await;
1362            tx.send(f(this, cx).await).ok();
1363        }));
1364
1365        cx.spawn(async move |this, cx| {
1366            let response = rx.await;
1367
1368            this.update(cx, |this, cx| this.update_last_checkpoint(cx))?
1369                .await?;
1370
1371            this.update(cx, |this, cx| {
1372                this.project
1373                    .update(cx, |project, cx| project.set_agent_location(None, cx));
1374                match response {
1375                    Ok(Err(e)) => {
1376                        this.send_task.take();
1377                        cx.emit(AcpThreadEvent::Error);
1378                        Err(e)
1379                    }
1380                    result => {
1381                        let canceled = matches!(
1382                            result,
1383                            Ok(Ok(acp::PromptResponse {
1384                                stop_reason: acp::StopReason::Canceled
1385                            }))
1386                        );
1387
1388                        // We only take the task if the current prompt wasn't canceled.
1389                        //
1390                        // This prompt may have been canceled because another one was sent
1391                        // while it was still generating. In these cases, dropping `send_task`
1392                        // would cause the next generation to be canceled.
1393                        if !canceled {
1394                            this.send_task.take();
1395                        }
1396
1397                        cx.emit(AcpThreadEvent::Stopped);
1398                        Ok(())
1399                    }
1400                }
1401            })?
1402        })
1403        .boxed()
1404    }
1405
1406    pub fn cancel(&mut self, cx: &mut Context<Self>) -> Task<()> {
1407        let Some(send_task) = self.send_task.take() else {
1408            return Task::ready(());
1409        };
1410
1411        for entry in self.entries.iter_mut() {
1412            if let AgentThreadEntry::ToolCall(call) = entry {
1413                let cancel = matches!(
1414                    call.status,
1415                    ToolCallStatus::Pending
1416                        | ToolCallStatus::WaitingForConfirmation { .. }
1417                        | ToolCallStatus::InProgress
1418                );
1419
1420                if cancel {
1421                    call.status = ToolCallStatus::Canceled;
1422                }
1423            }
1424        }
1425
1426        self.connection.cancel(&self.session_id, cx);
1427
1428        // Wait for the send task to complete
1429        cx.foreground_executor().spawn(send_task)
1430    }
1431
1432    /// Rewinds this thread to before the entry at `index`, removing it and all
1433    /// subsequent entries while reverting any changes made from that point.
1434    pub fn rewind(&mut self, id: UserMessageId, cx: &mut Context<Self>) -> Task<Result<()>> {
1435        let Some(session_editor) = self.connection.session_editor(&self.session_id, cx) else {
1436            return Task::ready(Err(anyhow!("not supported")));
1437        };
1438        let Some(message) = self.user_message(&id) else {
1439            return Task::ready(Err(anyhow!("message not found")));
1440        };
1441
1442        let checkpoint = message
1443            .checkpoint
1444            .as_ref()
1445            .map(|c| c.git_checkpoint.clone());
1446
1447        let git_store = self.project.read(cx).git_store().clone();
1448        cx.spawn(async move |this, cx| {
1449            if let Some(checkpoint) = checkpoint {
1450                git_store
1451                    .update(cx, |git, cx| git.restore_checkpoint(checkpoint, cx))?
1452                    .await?;
1453            }
1454
1455            cx.update(|cx| session_editor.truncate(id.clone(), cx))?
1456                .await?;
1457            this.update(cx, |this, cx| {
1458                if let Some((ix, _)) = this.user_message_mut(&id) {
1459                    let range = ix..this.entries.len();
1460                    this.entries.truncate(ix);
1461                    cx.emit(AcpThreadEvent::EntriesRemoved(range));
1462                }
1463            })
1464        })
1465    }
1466
1467    fn update_last_checkpoint(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
1468        let git_store = self.project.read(cx).git_store().clone();
1469
1470        let old_checkpoint = if let Some((_, message)) = self.last_user_message() {
1471            if let Some(checkpoint) = message.checkpoint.as_ref() {
1472                checkpoint.git_checkpoint.clone()
1473            } else {
1474                return Task::ready(Ok(()));
1475            }
1476        } else {
1477            return Task::ready(Ok(()));
1478        };
1479
1480        let new_checkpoint = git_store.update(cx, |git, cx| git.checkpoint(cx));
1481        cx.spawn(async move |this, cx| {
1482            let new_checkpoint = new_checkpoint
1483                .await
1484                .context("failed to get new checkpoint")
1485                .log_err();
1486            if let Some(new_checkpoint) = new_checkpoint {
1487                let equal = git_store
1488                    .update(cx, |git, cx| {
1489                        git.compare_checkpoints(old_checkpoint.clone(), new_checkpoint, cx)
1490                    })?
1491                    .await
1492                    .unwrap_or(true);
1493                this.update(cx, |this, cx| {
1494                    let (ix, message) = this.last_user_message().context("no user message")?;
1495                    let checkpoint = message.checkpoint.as_mut().context("no checkpoint")?;
1496                    checkpoint.show = !equal;
1497                    cx.emit(AcpThreadEvent::EntryUpdated(ix));
1498                    anyhow::Ok(())
1499                })??;
1500            }
1501
1502            Ok(())
1503        })
1504    }
1505
1506    fn last_user_message(&mut self) -> Option<(usize, &mut UserMessage)> {
1507        self.entries
1508            .iter_mut()
1509            .enumerate()
1510            .rev()
1511            .find_map(|(ix, entry)| {
1512                if let AgentThreadEntry::UserMessage(message) = entry {
1513                    Some((ix, message))
1514                } else {
1515                    None
1516                }
1517            })
1518    }
1519
1520    fn user_message(&self, id: &UserMessageId) -> Option<&UserMessage> {
1521        self.entries.iter().find_map(|entry| {
1522            if let AgentThreadEntry::UserMessage(message) = entry {
1523                if message.id.as_ref() == Some(id) {
1524                    Some(message)
1525                } else {
1526                    None
1527                }
1528            } else {
1529                None
1530            }
1531        })
1532    }
1533
1534    fn user_message_mut(&mut self, id: &UserMessageId) -> Option<(usize, &mut UserMessage)> {
1535        self.entries.iter_mut().enumerate().find_map(|(ix, entry)| {
1536            if let AgentThreadEntry::UserMessage(message) = entry {
1537                if message.id.as_ref() == Some(id) {
1538                    Some((ix, message))
1539                } else {
1540                    None
1541                }
1542            } else {
1543                None
1544            }
1545        })
1546    }
1547
1548    pub fn read_text_file(
1549        &self,
1550        path: PathBuf,
1551        line: Option<u32>,
1552        limit: Option<u32>,
1553        reuse_shared_snapshot: bool,
1554        cx: &mut Context<Self>,
1555    ) -> Task<Result<String>> {
1556        let project = self.project.clone();
1557        let action_log = self.action_log.clone();
1558        cx.spawn(async move |this, cx| {
1559            let load = project.update(cx, |project, cx| {
1560                let path = project
1561                    .project_path_for_absolute_path(&path, cx)
1562                    .context("invalid path")?;
1563                anyhow::Ok(project.open_buffer(path, cx))
1564            });
1565            let buffer = load??.await?;
1566
1567            let snapshot = if reuse_shared_snapshot {
1568                this.read_with(cx, |this, _| {
1569                    this.shared_buffers.get(&buffer.clone()).cloned()
1570                })
1571                .log_err()
1572                .flatten()
1573            } else {
1574                None
1575            };
1576
1577            let snapshot = if let Some(snapshot) = snapshot {
1578                snapshot
1579            } else {
1580                action_log.update(cx, |action_log, cx| {
1581                    action_log.buffer_read(buffer.clone(), cx);
1582                })?;
1583                project.update(cx, |project, cx| {
1584                    let position = buffer
1585                        .read(cx)
1586                        .snapshot()
1587                        .anchor_before(Point::new(line.unwrap_or_default(), 0));
1588                    project.set_agent_location(
1589                        Some(AgentLocation {
1590                            buffer: buffer.downgrade(),
1591                            position,
1592                        }),
1593                        cx,
1594                    );
1595                })?;
1596
1597                buffer.update(cx, |buffer, _| buffer.snapshot())?
1598            };
1599
1600            this.update(cx, |this, _| {
1601                let text = snapshot.text();
1602                this.shared_buffers.insert(buffer.clone(), snapshot);
1603                if line.is_none() && limit.is_none() {
1604                    return Ok(text);
1605                }
1606                let limit = limit.unwrap_or(u32::MAX) as usize;
1607                let Some(line) = line else {
1608                    return Ok(text.lines().take(limit).collect::<String>());
1609                };
1610
1611                let count = text.lines().count();
1612                if count < line as usize {
1613                    anyhow::bail!("There are only {} lines", count);
1614                }
1615                Ok(text
1616                    .lines()
1617                    .skip(line as usize + 1)
1618                    .take(limit)
1619                    .collect::<String>())
1620            })?
1621        })
1622    }
1623
1624    pub fn write_text_file(
1625        &self,
1626        path: PathBuf,
1627        content: String,
1628        cx: &mut Context<Self>,
1629    ) -> Task<Result<()>> {
1630        let project = self.project.clone();
1631        let action_log = self.action_log.clone();
1632        cx.spawn(async move |this, cx| {
1633            let load = project.update(cx, |project, cx| {
1634                let path = project
1635                    .project_path_for_absolute_path(&path, cx)
1636                    .context("invalid path")?;
1637                anyhow::Ok(project.open_buffer(path, cx))
1638            });
1639            let buffer = load??.await?;
1640            let snapshot = this.update(cx, |this, cx| {
1641                this.shared_buffers
1642                    .get(&buffer)
1643                    .cloned()
1644                    .unwrap_or_else(|| buffer.read(cx).snapshot())
1645            })?;
1646            let edits = cx
1647                .background_executor()
1648                .spawn(async move {
1649                    let old_text = snapshot.text();
1650                    text_diff(old_text.as_str(), &content)
1651                        .into_iter()
1652                        .map(|(range, replacement)| {
1653                            (
1654                                snapshot.anchor_after(range.start)
1655                                    ..snapshot.anchor_before(range.end),
1656                                replacement,
1657                            )
1658                        })
1659                        .collect::<Vec<_>>()
1660                })
1661                .await;
1662
1663            project.update(cx, |project, cx| {
1664                project.set_agent_location(
1665                    Some(AgentLocation {
1666                        buffer: buffer.downgrade(),
1667                        position: edits
1668                            .last()
1669                            .map(|(range, _)| range.end)
1670                            .unwrap_or(Anchor::MIN),
1671                    }),
1672                    cx,
1673                );
1674            })?;
1675
1676            let format_on_save = cx.update(|cx| {
1677                action_log.update(cx, |action_log, cx| {
1678                    action_log.buffer_read(buffer.clone(), cx);
1679                });
1680
1681                let format_on_save = buffer.update(cx, |buffer, cx| {
1682                    buffer.edit(edits, None, cx);
1683
1684                    let settings = language::language_settings::language_settings(
1685                        buffer.language().map(|l| l.name()),
1686                        buffer.file(),
1687                        cx,
1688                    );
1689
1690                    settings.format_on_save != FormatOnSave::Off
1691                });
1692                action_log.update(cx, |action_log, cx| {
1693                    action_log.buffer_edited(buffer.clone(), cx);
1694                });
1695                format_on_save
1696            })?;
1697
1698            if format_on_save {
1699                let format_task = project.update(cx, |project, cx| {
1700                    project.format(
1701                        HashSet::from_iter([buffer.clone()]),
1702                        LspFormatTarget::Buffers,
1703                        false,
1704                        FormatTrigger::Save,
1705                        cx,
1706                    )
1707                })?;
1708                format_task.await.log_err();
1709
1710                action_log.update(cx, |action_log, cx| {
1711                    action_log.buffer_edited(buffer.clone(), cx);
1712                })?;
1713            }
1714
1715            project
1716                .update(cx, |project, cx| project.save_buffer(buffer, cx))?
1717                .await
1718        })
1719    }
1720
1721    pub fn to_markdown(&self, cx: &App) -> String {
1722        self.entries.iter().map(|e| e.to_markdown(cx)).collect()
1723    }
1724
1725    pub fn emit_load_error(&mut self, error: LoadError, cx: &mut Context<Self>) {
1726        cx.emit(AcpThreadEvent::LoadError(error));
1727    }
1728}
1729
1730fn markdown_for_raw_output(
1731    raw_output: &serde_json::Value,
1732    language_registry: &Arc<LanguageRegistry>,
1733    cx: &mut App,
1734) -> Option<Entity<Markdown>> {
1735    match raw_output {
1736        serde_json::Value::Null => None,
1737        serde_json::Value::Bool(value) => Some(cx.new(|cx| {
1738            Markdown::new(
1739                value.to_string().into(),
1740                Some(language_registry.clone()),
1741                None,
1742                cx,
1743            )
1744        })),
1745        serde_json::Value::Number(value) => Some(cx.new(|cx| {
1746            Markdown::new(
1747                value.to_string().into(),
1748                Some(language_registry.clone()),
1749                None,
1750                cx,
1751            )
1752        })),
1753        serde_json::Value::String(value) => Some(cx.new(|cx| {
1754            Markdown::new(
1755                value.clone().into(),
1756                Some(language_registry.clone()),
1757                None,
1758                cx,
1759            )
1760        })),
1761        value => Some(cx.new(|cx| {
1762            Markdown::new(
1763                format!("```json\n{}\n```", value).into(),
1764                Some(language_registry.clone()),
1765                None,
1766                cx,
1767            )
1768        })),
1769    }
1770}
1771
1772#[cfg(test)]
1773mod tests {
1774    use super::*;
1775    use anyhow::anyhow;
1776    use futures::{channel::mpsc, future::LocalBoxFuture, select};
1777    use gpui::{App, AsyncApp, TestAppContext, WeakEntity};
1778    use indoc::indoc;
1779    use project::{FakeFs, Fs};
1780    use rand::Rng as _;
1781    use serde_json::json;
1782    use settings::SettingsStore;
1783    use smol::stream::StreamExt as _;
1784    use std::{
1785        any::Any,
1786        cell::RefCell,
1787        path::Path,
1788        rc::Rc,
1789        sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
1790        time::Duration,
1791    };
1792    use util::path;
1793
1794    fn init_test(cx: &mut TestAppContext) {
1795        env_logger::try_init().ok();
1796        cx.update(|cx| {
1797            let settings_store = SettingsStore::test(cx);
1798            cx.set_global(settings_store);
1799            Project::init_settings(cx);
1800            language::init(cx);
1801        });
1802    }
1803
1804    #[gpui::test]
1805    async fn test_push_user_content_block(cx: &mut gpui::TestAppContext) {
1806        init_test(cx);
1807
1808        let fs = FakeFs::new(cx.executor());
1809        let project = Project::test(fs, [], cx).await;
1810        let connection = Rc::new(FakeAgentConnection::new());
1811        let thread = cx
1812            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
1813            .await
1814            .unwrap();
1815
1816        // Test creating a new user message
1817        thread.update(cx, |thread, cx| {
1818            thread.push_user_content_block(
1819                None,
1820                acp::ContentBlock::Text(acp::TextContent {
1821                    annotations: None,
1822                    text: "Hello, ".to_string(),
1823                }),
1824                cx,
1825            );
1826        });
1827
1828        thread.update(cx, |thread, cx| {
1829            assert_eq!(thread.entries.len(), 1);
1830            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
1831                assert_eq!(user_msg.id, None);
1832                assert_eq!(user_msg.content.to_markdown(cx), "Hello, ");
1833            } else {
1834                panic!("Expected UserMessage");
1835            }
1836        });
1837
1838        // Test appending to existing user message
1839        let message_1_id = UserMessageId::new();
1840        thread.update(cx, |thread, cx| {
1841            thread.push_user_content_block(
1842                Some(message_1_id.clone()),
1843                acp::ContentBlock::Text(acp::TextContent {
1844                    annotations: None,
1845                    text: "world!".to_string(),
1846                }),
1847                cx,
1848            );
1849        });
1850
1851        thread.update(cx, |thread, cx| {
1852            assert_eq!(thread.entries.len(), 1);
1853            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
1854                assert_eq!(user_msg.id, Some(message_1_id));
1855                assert_eq!(user_msg.content.to_markdown(cx), "Hello, world!");
1856            } else {
1857                panic!("Expected UserMessage");
1858            }
1859        });
1860
1861        // Test creating new user message after assistant message
1862        thread.update(cx, |thread, cx| {
1863            thread.push_assistant_content_block(
1864                acp::ContentBlock::Text(acp::TextContent {
1865                    annotations: None,
1866                    text: "Assistant response".to_string(),
1867                }),
1868                false,
1869                cx,
1870            );
1871        });
1872
1873        let message_2_id = UserMessageId::new();
1874        thread.update(cx, |thread, cx| {
1875            thread.push_user_content_block(
1876                Some(message_2_id.clone()),
1877                acp::ContentBlock::Text(acp::TextContent {
1878                    annotations: None,
1879                    text: "New user message".to_string(),
1880                }),
1881                cx,
1882            );
1883        });
1884
1885        thread.update(cx, |thread, cx| {
1886            assert_eq!(thread.entries.len(), 3);
1887            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[2] {
1888                assert_eq!(user_msg.id, Some(message_2_id));
1889                assert_eq!(user_msg.content.to_markdown(cx), "New user message");
1890            } else {
1891                panic!("Expected UserMessage at index 2");
1892            }
1893        });
1894    }
1895
1896    #[gpui::test]
1897    async fn test_thinking_concatenation(cx: &mut gpui::TestAppContext) {
1898        init_test(cx);
1899
1900        let fs = FakeFs::new(cx.executor());
1901        let project = Project::test(fs, [], cx).await;
1902        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
1903            |_, thread, mut cx| {
1904                async move {
1905                    thread.update(&mut cx, |thread, cx| {
1906                        thread
1907                            .handle_session_update(
1908                                acp::SessionUpdate::AgentThoughtChunk {
1909                                    content: "Thinking ".into(),
1910                                },
1911                                cx,
1912                            )
1913                            .unwrap();
1914                        thread
1915                            .handle_session_update(
1916                                acp::SessionUpdate::AgentThoughtChunk {
1917                                    content: "hard!".into(),
1918                                },
1919                                cx,
1920                            )
1921                            .unwrap();
1922                    })?;
1923                    Ok(acp::PromptResponse {
1924                        stop_reason: acp::StopReason::EndTurn,
1925                    })
1926                }
1927                .boxed_local()
1928            },
1929        ));
1930
1931        let thread = cx
1932            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
1933            .await
1934            .unwrap();
1935
1936        thread
1937            .update(cx, |thread, cx| thread.send_raw("Hello from Zed!", cx))
1938            .await
1939            .unwrap();
1940
1941        let output = thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
1942        assert_eq!(
1943            output,
1944            indoc! {r#"
1945            ## User
1946
1947            Hello from Zed!
1948
1949            ## Assistant
1950
1951            <thinking>
1952            Thinking hard!
1953            </thinking>
1954
1955            "#}
1956        );
1957    }
1958
1959    #[gpui::test]
1960    async fn test_edits_concurrently_to_user(cx: &mut TestAppContext) {
1961        init_test(cx);
1962
1963        let fs = FakeFs::new(cx.executor());
1964        fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\n"}))
1965            .await;
1966        let project = Project::test(fs.clone(), [], cx).await;
1967        let (read_file_tx, read_file_rx) = oneshot::channel::<()>();
1968        let read_file_tx = Rc::new(RefCell::new(Some(read_file_tx)));
1969        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
1970            move |_, thread, mut cx| {
1971                let read_file_tx = read_file_tx.clone();
1972                async move {
1973                    let content = thread
1974                        .update(&mut cx, |thread, cx| {
1975                            thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
1976                        })
1977                        .unwrap()
1978                        .await
1979                        .unwrap();
1980                    assert_eq!(content, "one\ntwo\nthree\n");
1981                    read_file_tx.take().unwrap().send(()).unwrap();
1982                    thread
1983                        .update(&mut cx, |thread, cx| {
1984                            thread.write_text_file(
1985                                path!("/tmp/foo").into(),
1986                                "one\ntwo\nthree\nfour\nfive\n".to_string(),
1987                                cx,
1988                            )
1989                        })
1990                        .unwrap()
1991                        .await
1992                        .unwrap();
1993                    Ok(acp::PromptResponse {
1994                        stop_reason: acp::StopReason::EndTurn,
1995                    })
1996                }
1997                .boxed_local()
1998            },
1999        ));
2000
2001        let (worktree, pathbuf) = project
2002            .update(cx, |project, cx| {
2003                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
2004            })
2005            .await
2006            .unwrap();
2007        let buffer = project
2008            .update(cx, |project, cx| {
2009                project.open_buffer((worktree.read(cx).id(), pathbuf), cx)
2010            })
2011            .await
2012            .unwrap();
2013
2014        let thread = cx
2015            .update(|cx| connection.new_thread(project, Path::new(path!("/tmp")), cx))
2016            .await
2017            .unwrap();
2018
2019        let request = thread.update(cx, |thread, cx| {
2020            thread.send_raw("Extend the count in /tmp/foo", cx)
2021        });
2022        read_file_rx.await.ok();
2023        buffer.update(cx, |buffer, cx| {
2024            buffer.edit([(0..0, "zero\n".to_string())], None, cx);
2025        });
2026        cx.run_until_parked();
2027        assert_eq!(
2028            buffer.read_with(cx, |buffer, _| buffer.text()),
2029            "zero\none\ntwo\nthree\nfour\nfive\n"
2030        );
2031        assert_eq!(
2032            String::from_utf8(fs.read_file_sync(path!("/tmp/foo")).unwrap()).unwrap(),
2033            "zero\none\ntwo\nthree\nfour\nfive\n"
2034        );
2035        request.await.unwrap();
2036    }
2037
2038    #[gpui::test]
2039    async fn test_succeeding_canceled_toolcall(cx: &mut TestAppContext) {
2040        init_test(cx);
2041
2042        let fs = FakeFs::new(cx.executor());
2043        let project = Project::test(fs, [], cx).await;
2044        let id = acp::ToolCallId("test".into());
2045
2046        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2047            let id = id.clone();
2048            move |_, thread, mut cx| {
2049                let id = id.clone();
2050                async move {
2051                    thread
2052                        .update(&mut cx, |thread, cx| {
2053                            thread.handle_session_update(
2054                                acp::SessionUpdate::ToolCall(acp::ToolCall {
2055                                    id: id.clone(),
2056                                    title: "Label".into(),
2057                                    kind: acp::ToolKind::Fetch,
2058                                    status: acp::ToolCallStatus::InProgress,
2059                                    content: vec![],
2060                                    locations: vec![],
2061                                    raw_input: None,
2062                                    raw_output: None,
2063                                }),
2064                                cx,
2065                            )
2066                        })
2067                        .unwrap()
2068                        .unwrap();
2069                    Ok(acp::PromptResponse {
2070                        stop_reason: acp::StopReason::EndTurn,
2071                    })
2072                }
2073                .boxed_local()
2074            }
2075        }));
2076
2077        let thread = cx
2078            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2079            .await
2080            .unwrap();
2081
2082        let request = thread.update(cx, |thread, cx| {
2083            thread.send_raw("Fetch https://example.com", cx)
2084        });
2085
2086        run_until_first_tool_call(&thread, cx).await;
2087
2088        thread.read_with(cx, |thread, _| {
2089            assert!(matches!(
2090                thread.entries[1],
2091                AgentThreadEntry::ToolCall(ToolCall {
2092                    status: ToolCallStatus::InProgress,
2093                    ..
2094                })
2095            ));
2096        });
2097
2098        thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2099
2100        thread.read_with(cx, |thread, _| {
2101            assert!(matches!(
2102                &thread.entries[1],
2103                AgentThreadEntry::ToolCall(ToolCall {
2104                    status: ToolCallStatus::Canceled,
2105                    ..
2106                })
2107            ));
2108        });
2109
2110        thread
2111            .update(cx, |thread, cx| {
2112                thread.handle_session_update(
2113                    acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate {
2114                        id,
2115                        fields: acp::ToolCallUpdateFields {
2116                            status: Some(acp::ToolCallStatus::Completed),
2117                            ..Default::default()
2118                        },
2119                    }),
2120                    cx,
2121                )
2122            })
2123            .unwrap();
2124
2125        request.await.unwrap();
2126
2127        thread.read_with(cx, |thread, _| {
2128            assert!(matches!(
2129                thread.entries[1],
2130                AgentThreadEntry::ToolCall(ToolCall {
2131                    status: ToolCallStatus::Completed,
2132                    ..
2133                })
2134            ));
2135        });
2136    }
2137
2138    #[gpui::test]
2139    async fn test_no_pending_edits_if_tool_calls_are_completed(cx: &mut TestAppContext) {
2140        init_test(cx);
2141        let fs = FakeFs::new(cx.background_executor.clone());
2142        fs.insert_tree(path!("/test"), json!({})).await;
2143        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
2144
2145        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2146            move |_, thread, mut cx| {
2147                async move {
2148                    thread
2149                        .update(&mut cx, |thread, cx| {
2150                            thread.handle_session_update(
2151                                acp::SessionUpdate::ToolCall(acp::ToolCall {
2152                                    id: acp::ToolCallId("test".into()),
2153                                    title: "Label".into(),
2154                                    kind: acp::ToolKind::Edit,
2155                                    status: acp::ToolCallStatus::Completed,
2156                                    content: vec![acp::ToolCallContent::Diff {
2157                                        diff: acp::Diff {
2158                                            path: "/test/test.txt".into(),
2159                                            old_text: None,
2160                                            new_text: "foo".into(),
2161                                        },
2162                                    }],
2163                                    locations: vec![],
2164                                    raw_input: None,
2165                                    raw_output: None,
2166                                }),
2167                                cx,
2168                            )
2169                        })
2170                        .unwrap()
2171                        .unwrap();
2172                    Ok(acp::PromptResponse {
2173                        stop_reason: acp::StopReason::EndTurn,
2174                    })
2175                }
2176                .boxed_local()
2177            }
2178        }));
2179
2180        let thread = cx
2181            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2182            .await
2183            .unwrap();
2184
2185        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Hi".into()], cx)))
2186            .await
2187            .unwrap();
2188
2189        assert!(cx.read(|cx| !thread.read(cx).has_pending_edit_tool_calls()));
2190    }
2191
2192    #[gpui::test(iterations = 10)]
2193    async fn test_checkpoints(cx: &mut TestAppContext) {
2194        init_test(cx);
2195        let fs = FakeFs::new(cx.background_executor.clone());
2196        fs.insert_tree(
2197            path!("/test"),
2198            json!({
2199                ".git": {}
2200            }),
2201        )
2202        .await;
2203        let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
2204
2205        let simulate_changes = Arc::new(AtomicBool::new(true));
2206        let next_filename = Arc::new(AtomicUsize::new(0));
2207        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
2208            let simulate_changes = simulate_changes.clone();
2209            let next_filename = next_filename.clone();
2210            let fs = fs.clone();
2211            move |request, thread, mut cx| {
2212                let fs = fs.clone();
2213                let simulate_changes = simulate_changes.clone();
2214                let next_filename = next_filename.clone();
2215                async move {
2216                    if simulate_changes.load(SeqCst) {
2217                        let filename = format!("/test/file-{}", next_filename.fetch_add(1, SeqCst));
2218                        fs.write(Path::new(&filename), b"").await?;
2219                    }
2220
2221                    let acp::ContentBlock::Text(content) = &request.prompt[0] else {
2222                        panic!("expected text content block");
2223                    };
2224                    thread.update(&mut cx, |thread, cx| {
2225                        thread
2226                            .handle_session_update(
2227                                acp::SessionUpdate::AgentMessageChunk {
2228                                    content: content.text.to_uppercase().into(),
2229                                },
2230                                cx,
2231                            )
2232                            .unwrap();
2233                    })?;
2234                    Ok(acp::PromptResponse {
2235                        stop_reason: acp::StopReason::EndTurn,
2236                    })
2237                }
2238                .boxed_local()
2239            }
2240        }));
2241        let thread = cx
2242            .update(|cx| connection.new_thread(project, Path::new(path!("/test")), cx))
2243            .await
2244            .unwrap();
2245
2246        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Lorem".into()], cx)))
2247            .await
2248            .unwrap();
2249        thread.read_with(cx, |thread, cx| {
2250            assert_eq!(
2251                thread.to_markdown(cx),
2252                indoc! {"
2253                    ## User (checkpoint)
2254
2255                    Lorem
2256
2257                    ## Assistant
2258
2259                    LOREM
2260
2261                "}
2262            );
2263        });
2264        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2265
2266        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["ipsum".into()], cx)))
2267            .await
2268            .unwrap();
2269        thread.read_with(cx, |thread, cx| {
2270            assert_eq!(
2271                thread.to_markdown(cx),
2272                indoc! {"
2273                    ## User (checkpoint)
2274
2275                    Lorem
2276
2277                    ## Assistant
2278
2279                    LOREM
2280
2281                    ## User (checkpoint)
2282
2283                    ipsum
2284
2285                    ## Assistant
2286
2287                    IPSUM
2288
2289                "}
2290            );
2291        });
2292        assert_eq!(
2293            fs.files(),
2294            vec![
2295                Path::new(path!("/test/file-0")),
2296                Path::new(path!("/test/file-1"))
2297            ]
2298        );
2299
2300        // Checkpoint isn't stored when there are no changes.
2301        simulate_changes.store(false, SeqCst);
2302        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["dolor".into()], cx)))
2303            .await
2304            .unwrap();
2305        thread.read_with(cx, |thread, cx| {
2306            assert_eq!(
2307                thread.to_markdown(cx),
2308                indoc! {"
2309                    ## User (checkpoint)
2310
2311                    Lorem
2312
2313                    ## Assistant
2314
2315                    LOREM
2316
2317                    ## User (checkpoint)
2318
2319                    ipsum
2320
2321                    ## Assistant
2322
2323                    IPSUM
2324
2325                    ## User
2326
2327                    dolor
2328
2329                    ## Assistant
2330
2331                    DOLOR
2332
2333                "}
2334            );
2335        });
2336        assert_eq!(
2337            fs.files(),
2338            vec![
2339                Path::new(path!("/test/file-0")),
2340                Path::new(path!("/test/file-1"))
2341            ]
2342        );
2343
2344        // Rewinding the conversation truncates the history and restores the checkpoint.
2345        thread
2346            .update(cx, |thread, cx| {
2347                let AgentThreadEntry::UserMessage(message) = &thread.entries[2] else {
2348                    panic!("unexpected entries {:?}", thread.entries)
2349                };
2350                thread.rewind(message.id.clone().unwrap(), cx)
2351            })
2352            .await
2353            .unwrap();
2354        thread.read_with(cx, |thread, cx| {
2355            assert_eq!(
2356                thread.to_markdown(cx),
2357                indoc! {"
2358                    ## User (checkpoint)
2359
2360                    Lorem
2361
2362                    ## Assistant
2363
2364                    LOREM
2365
2366                "}
2367            );
2368        });
2369        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
2370    }
2371
2372    async fn run_until_first_tool_call(
2373        thread: &Entity<AcpThread>,
2374        cx: &mut TestAppContext,
2375    ) -> usize {
2376        let (mut tx, mut rx) = mpsc::channel::<usize>(1);
2377
2378        let subscription = cx.update(|cx| {
2379            cx.subscribe(thread, move |thread, _, cx| {
2380                for (ix, entry) in thread.read(cx).entries.iter().enumerate() {
2381                    if matches!(entry, AgentThreadEntry::ToolCall(_)) {
2382                        return tx.try_send(ix).unwrap();
2383                    }
2384                }
2385            })
2386        });
2387
2388        select! {
2389            _ = futures::FutureExt::fuse(smol::Timer::after(Duration::from_secs(10))) => {
2390                panic!("Timeout waiting for tool call")
2391            }
2392            ix = rx.next().fuse() => {
2393                drop(subscription);
2394                ix.unwrap()
2395            }
2396        }
2397    }
2398
2399    #[derive(Clone, Default)]
2400    struct FakeAgentConnection {
2401        auth_methods: Vec<acp::AuthMethod>,
2402        sessions: Arc<parking_lot::Mutex<HashMap<acp::SessionId, WeakEntity<AcpThread>>>>,
2403        on_user_message: Option<
2404            Rc<
2405                dyn Fn(
2406                        acp::PromptRequest,
2407                        WeakEntity<AcpThread>,
2408                        AsyncApp,
2409                    ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
2410                    + 'static,
2411            >,
2412        >,
2413    }
2414
2415    impl FakeAgentConnection {
2416        fn new() -> Self {
2417            Self {
2418                auth_methods: Vec::new(),
2419                on_user_message: None,
2420                sessions: Arc::default(),
2421            }
2422        }
2423
2424        #[expect(unused)]
2425        fn with_auth_methods(mut self, auth_methods: Vec<acp::AuthMethod>) -> Self {
2426            self.auth_methods = auth_methods;
2427            self
2428        }
2429
2430        fn on_user_message(
2431            mut self,
2432            handler: impl Fn(
2433                acp::PromptRequest,
2434                WeakEntity<AcpThread>,
2435                AsyncApp,
2436            ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
2437            + 'static,
2438        ) -> Self {
2439            self.on_user_message.replace(Rc::new(handler));
2440            self
2441        }
2442    }
2443
2444    impl AgentConnection for FakeAgentConnection {
2445        fn auth_methods(&self) -> &[acp::AuthMethod] {
2446            &self.auth_methods
2447        }
2448
2449        fn new_thread(
2450            self: Rc<Self>,
2451            project: Entity<Project>,
2452            _cwd: &Path,
2453            cx: &mut App,
2454        ) -> Task<gpui::Result<Entity<AcpThread>>> {
2455            let session_id = acp::SessionId(
2456                rand::thread_rng()
2457                    .sample_iter(&rand::distributions::Alphanumeric)
2458                    .take(7)
2459                    .map(char::from)
2460                    .collect::<String>()
2461                    .into(),
2462            );
2463            let action_log = cx.new(|_| ActionLog::new(project.clone()));
2464            let thread = cx.new(|_cx| {
2465                AcpThread::new(
2466                    "Test",
2467                    self.clone(),
2468                    project,
2469                    action_log,
2470                    session_id.clone(),
2471                )
2472            });
2473            self.sessions.lock().insert(session_id, thread.downgrade());
2474            Task::ready(Ok(thread))
2475        }
2476
2477        fn authenticate(&self, method: acp::AuthMethodId, _cx: &mut App) -> Task<gpui::Result<()>> {
2478            if self.auth_methods().iter().any(|m| m.id == method) {
2479                Task::ready(Ok(()))
2480            } else {
2481                Task::ready(Err(anyhow!("Invalid Auth Method")))
2482            }
2483        }
2484
2485        fn prompt(
2486            &self,
2487            _id: Option<UserMessageId>,
2488            params: acp::PromptRequest,
2489            cx: &mut App,
2490        ) -> Task<gpui::Result<acp::PromptResponse>> {
2491            let sessions = self.sessions.lock();
2492            let thread = sessions.get(&params.session_id).unwrap();
2493            if let Some(handler) = &self.on_user_message {
2494                let handler = handler.clone();
2495                let thread = thread.clone();
2496                cx.spawn(async move |cx| handler(params, thread, cx.clone()).await)
2497            } else {
2498                Task::ready(Ok(acp::PromptResponse {
2499                    stop_reason: acp::StopReason::EndTurn,
2500                }))
2501            }
2502        }
2503
2504        fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
2505            let sessions = self.sessions.lock();
2506            let thread = sessions.get(session_id).unwrap().clone();
2507
2508            cx.spawn(async move |cx| {
2509                thread
2510                    .update(cx, |thread, cx| thread.cancel(cx))
2511                    .unwrap()
2512                    .await
2513            })
2514            .detach();
2515        }
2516
2517        fn session_editor(
2518            &self,
2519            session_id: &acp::SessionId,
2520            _cx: &mut App,
2521        ) -> Option<Rc<dyn AgentSessionEditor>> {
2522            Some(Rc::new(FakeAgentSessionEditor {
2523                _session_id: session_id.clone(),
2524            }))
2525        }
2526
2527        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
2528            self
2529        }
2530    }
2531
2532    struct FakeAgentSessionEditor {
2533        _session_id: acp::SessionId,
2534    }
2535
2536    impl AgentSessionEditor for FakeAgentSessionEditor {
2537        fn truncate(&self, _message_id: UserMessageId, _cx: &mut App) -> Task<Result<()>> {
2538            Task::ready(Ok(()))
2539        }
2540    }
2541}