acp_thread.rs

   1mod connection;
   2mod diff;
   3mod mention;
   4mod terminal;
   5use action_log::{ActionLog, ActionLogTelemetry};
   6use agent_client_protocol::{self as acp};
   7use anyhow::{Context as _, Result, anyhow};
   8use collections::HashSet;
   9pub use connection::*;
  10pub use diff::*;
  11use futures::{FutureExt, channel::oneshot, future::BoxFuture};
  12use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
  13use itertools::Itertools;
  14use language::language_settings::FormatOnSave;
  15use language::{Anchor, Buffer, BufferSnapshot, LanguageRegistry, Point, ToPoint, text_diff};
  16use markdown::Markdown;
  17pub use mention::*;
  18use project::lsp_store::{FormatTrigger, LspFormatTarget};
  19use project::{AgentLocation, Project, git_store::GitStoreCheckpoint};
  20use serde::{Deserialize, Serialize};
  21use serde_json::to_string_pretty;
  22use std::collections::HashMap;
  23use std::error::Error;
  24use std::fmt::{Formatter, Write};
  25use std::ops::Range;
  26use std::process::ExitStatus;
  27use std::rc::Rc;
  28use std::time::{Duration, Instant};
  29use std::{fmt::Display, mem, path::PathBuf, sync::Arc};
  30use task::{Shell, ShellBuilder};
  31pub use terminal::*;
  32use text::Bias;
  33use ui::App;
  34use util::{ResultExt, get_default_system_shell_preferring_bash, paths::PathStyle};
  35use uuid::Uuid;
  36
  37/// Key used in ACP ToolCall meta to store the tool's programmatic name.
  38/// This is a workaround since ACP's ToolCall doesn't have a dedicated name field.
  39pub const TOOL_NAME_META_KEY: &str = "tool_name";
  40
  41/// Helper to extract tool name from ACP meta
  42pub fn tool_name_from_meta(meta: &Option<acp::Meta>) -> Option<SharedString> {
  43    meta.as_ref()
  44        .and_then(|m| m.get(TOOL_NAME_META_KEY))
  45        .and_then(|v| v.as_str())
  46        .map(|s| SharedString::from(s.to_owned()))
  47}
  48
  49/// Helper to create meta with tool name
  50pub fn meta_with_tool_name(tool_name: &str) -> acp::Meta {
  51    acp::Meta::from_iter([(TOOL_NAME_META_KEY.into(), tool_name.into())])
  52}
  53
  54/// Key used in ACP ToolCall meta to store the session id and message indexes
  55pub const SUBAGENT_SESSION_INFO_META_KEY: &str = "subagent_session_info";
  56
  57#[derive(Clone, Debug, Deserialize, Serialize)]
  58pub struct SubagentSessionInfo {
  59    /// The session id of the subagent sessiont that was spawned
  60    pub session_id: acp::SessionId,
  61    /// The index of the message of the start of the "turn" run by this tool call
  62    pub message_start_index: usize,
  63    /// The index of the output of the message that the subagent has returned
  64    #[serde(skip_serializing_if = "Option::is_none")]
  65    pub message_end_index: Option<usize>,
  66}
  67
  68/// Helper to extract subagent session id from ACP meta
  69pub fn subagent_session_info_from_meta(meta: &Option<acp::Meta>) -> Option<SubagentSessionInfo> {
  70    meta.as_ref()
  71        .and_then(|m| m.get(SUBAGENT_SESSION_INFO_META_KEY))
  72        .and_then(|v| serde_json::from_value(v.clone()).ok())
  73}
  74
  75#[derive(Debug)]
  76pub struct UserMessage {
  77    pub id: Option<UserMessageId>,
  78    pub content: ContentBlock,
  79    pub chunks: Vec<acp::ContentBlock>,
  80    pub checkpoint: Option<Checkpoint>,
  81    pub indented: bool,
  82}
  83
  84#[derive(Debug)]
  85pub struct Checkpoint {
  86    git_checkpoint: GitStoreCheckpoint,
  87    pub show: bool,
  88}
  89
  90impl UserMessage {
  91    fn to_markdown(&self, cx: &App) -> String {
  92        let mut markdown = String::new();
  93        if self
  94            .checkpoint
  95            .as_ref()
  96            .is_some_and(|checkpoint| checkpoint.show)
  97        {
  98            writeln!(markdown, "## User (checkpoint)").unwrap();
  99        } else {
 100            writeln!(markdown, "## User").unwrap();
 101        }
 102        writeln!(markdown).unwrap();
 103        writeln!(markdown, "{}", self.content.to_markdown(cx)).unwrap();
 104        writeln!(markdown).unwrap();
 105        markdown
 106    }
 107}
 108
 109#[derive(Debug, PartialEq)]
 110pub struct AssistantMessage {
 111    pub chunks: Vec<AssistantMessageChunk>,
 112    pub indented: bool,
 113    pub is_subagent_output: bool,
 114}
 115
 116impl AssistantMessage {
 117    pub fn to_markdown(&self, cx: &App) -> String {
 118        format!(
 119            "## Assistant\n\n{}\n\n",
 120            self.chunks
 121                .iter()
 122                .map(|chunk| chunk.to_markdown(cx))
 123                .join("\n\n")
 124        )
 125    }
 126}
 127
 128#[derive(Debug, PartialEq)]
 129pub enum AssistantMessageChunk {
 130    Message { block: ContentBlock },
 131    Thought { block: ContentBlock },
 132}
 133
 134impl AssistantMessageChunk {
 135    pub fn from_str(
 136        chunk: &str,
 137        language_registry: &Arc<LanguageRegistry>,
 138        path_style: PathStyle,
 139        cx: &mut App,
 140    ) -> Self {
 141        Self::Message {
 142            block: ContentBlock::new(chunk.into(), language_registry, path_style, cx),
 143        }
 144    }
 145
 146    fn to_markdown(&self, cx: &App) -> String {
 147        match self {
 148            Self::Message { block } => block.to_markdown(cx).to_string(),
 149            Self::Thought { block } => {
 150                format!("<thinking>\n{}\n</thinking>", block.to_markdown(cx))
 151            }
 152        }
 153    }
 154}
 155
 156#[derive(Debug)]
 157pub enum AgentThreadEntry {
 158    UserMessage(UserMessage),
 159    AssistantMessage(AssistantMessage),
 160    ToolCall(ToolCall),
 161}
 162
 163impl AgentThreadEntry {
 164    pub fn is_indented(&self) -> bool {
 165        match self {
 166            Self::UserMessage(message) => message.indented,
 167            Self::AssistantMessage(message) => message.indented,
 168            Self::ToolCall(_) => false,
 169        }
 170    }
 171
 172    pub fn to_markdown(&self, cx: &App) -> String {
 173        match self {
 174            Self::UserMessage(message) => message.to_markdown(cx),
 175            Self::AssistantMessage(message) => message.to_markdown(cx),
 176            Self::ToolCall(tool_call) => tool_call.to_markdown(cx),
 177        }
 178    }
 179
 180    pub fn user_message(&self) -> Option<&UserMessage> {
 181        if let AgentThreadEntry::UserMessage(message) = self {
 182            Some(message)
 183        } else {
 184            None
 185        }
 186    }
 187
 188    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 189        if let AgentThreadEntry::ToolCall(call) = self {
 190            itertools::Either::Left(call.diffs())
 191        } else {
 192            itertools::Either::Right(std::iter::empty())
 193        }
 194    }
 195
 196    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 197        if let AgentThreadEntry::ToolCall(call) = self {
 198            itertools::Either::Left(call.terminals())
 199        } else {
 200            itertools::Either::Right(std::iter::empty())
 201        }
 202    }
 203
 204    pub fn location(&self, ix: usize) -> Option<(acp::ToolCallLocation, AgentLocation)> {
 205        if let AgentThreadEntry::ToolCall(ToolCall {
 206            locations,
 207            resolved_locations,
 208            ..
 209        }) = self
 210        {
 211            Some((
 212                locations.get(ix)?.clone(),
 213                resolved_locations.get(ix)?.clone()?,
 214            ))
 215        } else {
 216            None
 217        }
 218    }
 219}
 220
 221#[derive(Debug)]
 222pub struct ToolCall {
 223    pub id: acp::ToolCallId,
 224    pub label: Entity<Markdown>,
 225    pub kind: acp::ToolKind,
 226    pub content: Vec<ToolCallContent>,
 227    pub status: ToolCallStatus,
 228    pub locations: Vec<acp::ToolCallLocation>,
 229    pub resolved_locations: Vec<Option<AgentLocation>>,
 230    pub raw_input: Option<serde_json::Value>,
 231    pub raw_input_markdown: Option<Entity<Markdown>>,
 232    pub raw_output: Option<serde_json::Value>,
 233    pub tool_name: Option<SharedString>,
 234    pub subagent_session_info: Option<SubagentSessionInfo>,
 235}
 236
 237impl ToolCall {
 238    fn from_acp(
 239        tool_call: acp::ToolCall,
 240        status: ToolCallStatus,
 241        language_registry: Arc<LanguageRegistry>,
 242        path_style: PathStyle,
 243        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 244        cx: &mut App,
 245    ) -> Result<Self> {
 246        let title = if tool_call.kind == acp::ToolKind::Execute {
 247            tool_call.title
 248        } else if let Some((first_line, _)) = tool_call.title.split_once("\n") {
 249            first_line.to_owned() + ""
 250        } else {
 251            tool_call.title
 252        };
 253        let mut content = Vec::with_capacity(tool_call.content.len());
 254        for item in tool_call.content {
 255            if let Some(item) = ToolCallContent::from_acp(
 256                item,
 257                language_registry.clone(),
 258                path_style,
 259                terminals,
 260                cx,
 261            )? {
 262                content.push(item);
 263            }
 264        }
 265
 266        let raw_input_markdown = tool_call
 267            .raw_input
 268            .as_ref()
 269            .and_then(|input| markdown_for_raw_output(input, &language_registry, cx));
 270
 271        let tool_name = tool_name_from_meta(&tool_call.meta);
 272
 273        let subagent_session_info = subagent_session_info_from_meta(&tool_call.meta);
 274
 275        let result = Self {
 276            id: tool_call.tool_call_id,
 277            label: cx
 278                .new(|cx| Markdown::new(title.into(), Some(language_registry.clone()), None, cx)),
 279            kind: tool_call.kind,
 280            content,
 281            locations: tool_call.locations,
 282            resolved_locations: Vec::default(),
 283            status,
 284            raw_input: tool_call.raw_input,
 285            raw_input_markdown,
 286            raw_output: tool_call.raw_output,
 287            tool_name,
 288            subagent_session_info,
 289        };
 290        Ok(result)
 291    }
 292
 293    fn update_fields(
 294        &mut self,
 295        fields: acp::ToolCallUpdateFields,
 296        meta: Option<acp::Meta>,
 297        language_registry: Arc<LanguageRegistry>,
 298        path_style: PathStyle,
 299        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 300        cx: &mut App,
 301    ) -> Result<()> {
 302        let acp::ToolCallUpdateFields {
 303            kind,
 304            status,
 305            title,
 306            content,
 307            locations,
 308            raw_input,
 309            raw_output,
 310            ..
 311        } = fields;
 312
 313        if let Some(kind) = kind {
 314            self.kind = kind;
 315        }
 316
 317        if let Some(status) = status {
 318            self.status = status.into();
 319        }
 320
 321        if let Some(subagent_session_info) = subagent_session_info_from_meta(&meta) {
 322            self.subagent_session_info = Some(subagent_session_info);
 323        }
 324
 325        if let Some(title) = title {
 326            if self.kind == acp::ToolKind::Execute {
 327                for terminal in self.terminals() {
 328                    terminal.update(cx, |terminal, cx| {
 329                        terminal.update_command_label(&title, cx);
 330                    });
 331                }
 332            }
 333            self.label.update(cx, |label, cx| {
 334                if self.kind == acp::ToolKind::Execute {
 335                    label.replace(title, cx);
 336                } else if let Some((first_line, _)) = title.split_once("\n") {
 337                    label.replace(first_line.to_owned() + "", cx);
 338                } else {
 339                    label.replace(title, cx);
 340                }
 341            });
 342        }
 343
 344        if let Some(content) = content {
 345            let mut new_content_len = content.len();
 346            let mut content = content.into_iter();
 347
 348            // Reuse existing content if we can
 349            for (old, new) in self.content.iter_mut().zip(content.by_ref()) {
 350                let valid_content =
 351                    old.update_from_acp(new, language_registry.clone(), path_style, terminals, cx)?;
 352                if !valid_content {
 353                    new_content_len -= 1;
 354                }
 355            }
 356            for new in content {
 357                if let Some(new) = ToolCallContent::from_acp(
 358                    new,
 359                    language_registry.clone(),
 360                    path_style,
 361                    terminals,
 362                    cx,
 363                )? {
 364                    self.content.push(new);
 365                } else {
 366                    new_content_len -= 1;
 367                }
 368            }
 369            self.content.truncate(new_content_len);
 370        }
 371
 372        if let Some(locations) = locations {
 373            self.locations = locations;
 374        }
 375
 376        if let Some(raw_input) = raw_input {
 377            self.raw_input_markdown = markdown_for_raw_output(&raw_input, &language_registry, cx);
 378            self.raw_input = Some(raw_input);
 379        }
 380
 381        if let Some(raw_output) = raw_output {
 382            if self.content.is_empty()
 383                && let Some(markdown) = markdown_for_raw_output(&raw_output, &language_registry, cx)
 384            {
 385                self.content
 386                    .push(ToolCallContent::ContentBlock(ContentBlock::Markdown {
 387                        markdown,
 388                    }));
 389            }
 390            self.raw_output = Some(raw_output);
 391        }
 392        Ok(())
 393    }
 394
 395    pub fn diffs(&self) -> impl Iterator<Item = &Entity<Diff>> {
 396        self.content.iter().filter_map(|content| match content {
 397            ToolCallContent::Diff(diff) => Some(diff),
 398            ToolCallContent::ContentBlock(_) => None,
 399            ToolCallContent::Terminal(_) => None,
 400        })
 401    }
 402
 403    pub fn terminals(&self) -> impl Iterator<Item = &Entity<Terminal>> {
 404        self.content.iter().filter_map(|content| match content {
 405            ToolCallContent::Terminal(terminal) => Some(terminal),
 406            ToolCallContent::ContentBlock(_) => None,
 407            ToolCallContent::Diff(_) => None,
 408        })
 409    }
 410
 411    pub fn is_subagent(&self) -> bool {
 412        self.tool_name.as_ref().is_some_and(|s| s == "spawn_agent")
 413            || self.subagent_session_info.is_some()
 414    }
 415
 416    pub fn to_markdown(&self, cx: &App) -> String {
 417        let mut markdown = format!(
 418            "**Tool Call: {}**\nStatus: {}\n\n",
 419            self.label.read(cx).source(),
 420            self.status
 421        );
 422        for content in &self.content {
 423            markdown.push_str(content.to_markdown(cx).as_str());
 424            markdown.push_str("\n\n");
 425        }
 426        markdown
 427    }
 428
 429    async fn resolve_location(
 430        location: acp::ToolCallLocation,
 431        project: WeakEntity<Project>,
 432        cx: &mut AsyncApp,
 433    ) -> Option<ResolvedLocation> {
 434        let buffer = project
 435            .update(cx, |project, cx| {
 436                project
 437                    .project_path_for_absolute_path(&location.path, cx)
 438                    .map(|path| project.open_buffer(path, cx))
 439            })
 440            .ok()??;
 441        let buffer = buffer.await.log_err()?;
 442        let position = buffer.update(cx, |buffer, _| {
 443            let snapshot = buffer.snapshot();
 444            if let Some(row) = location.line {
 445                let column = snapshot.indent_size_for_line(row).len;
 446                let point = snapshot.clip_point(Point::new(row, column), Bias::Left);
 447                snapshot.anchor_before(point)
 448            } else {
 449                Anchor::min_for_buffer(snapshot.remote_id())
 450            }
 451        });
 452
 453        Some(ResolvedLocation { buffer, position })
 454    }
 455
 456    fn resolve_locations(
 457        &self,
 458        project: Entity<Project>,
 459        cx: &mut App,
 460    ) -> Task<Vec<Option<ResolvedLocation>>> {
 461        let locations = self.locations.clone();
 462        project.update(cx, |_, cx| {
 463            cx.spawn(async move |project, cx| {
 464                let mut new_locations = Vec::new();
 465                for location in locations {
 466                    new_locations.push(Self::resolve_location(location, project.clone(), cx).await);
 467                }
 468                new_locations
 469            })
 470        })
 471    }
 472}
 473
 474// Separate so we can hold a strong reference to the buffer
 475// for saving on the thread
 476#[derive(Clone, Debug, PartialEq, Eq)]
 477struct ResolvedLocation {
 478    buffer: Entity<Buffer>,
 479    position: Anchor,
 480}
 481
 482impl From<&ResolvedLocation> for AgentLocation {
 483    fn from(value: &ResolvedLocation) -> Self {
 484        Self {
 485            buffer: value.buffer.downgrade(),
 486            position: value.position,
 487        }
 488    }
 489}
 490
 491#[derive(Debug)]
 492pub enum ToolCallStatus {
 493    /// The tool call hasn't started running yet, but we start showing it to
 494    /// the user.
 495    Pending,
 496    /// The tool call is waiting for confirmation from the user.
 497    WaitingForConfirmation {
 498        options: PermissionOptions,
 499        respond_tx: oneshot::Sender<acp::PermissionOptionId>,
 500    },
 501    /// The tool call is currently running.
 502    InProgress,
 503    /// The tool call completed successfully.
 504    Completed,
 505    /// The tool call failed.
 506    Failed,
 507    /// The user rejected the tool call.
 508    Rejected,
 509    /// The user canceled generation so the tool call was canceled.
 510    Canceled,
 511}
 512
 513impl From<acp::ToolCallStatus> for ToolCallStatus {
 514    fn from(status: acp::ToolCallStatus) -> Self {
 515        match status {
 516            acp::ToolCallStatus::Pending => Self::Pending,
 517            acp::ToolCallStatus::InProgress => Self::InProgress,
 518            acp::ToolCallStatus::Completed => Self::Completed,
 519            acp::ToolCallStatus::Failed => Self::Failed,
 520            _ => Self::Pending,
 521        }
 522    }
 523}
 524
 525impl Display for ToolCallStatus {
 526    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
 527        write!(
 528            f,
 529            "{}",
 530            match self {
 531                ToolCallStatus::Pending => "Pending",
 532                ToolCallStatus::WaitingForConfirmation { .. } => "Waiting for confirmation",
 533                ToolCallStatus::InProgress => "In Progress",
 534                ToolCallStatus::Completed => "Completed",
 535                ToolCallStatus::Failed => "Failed",
 536                ToolCallStatus::Rejected => "Rejected",
 537                ToolCallStatus::Canceled => "Canceled",
 538            }
 539        )
 540    }
 541}
 542
 543#[derive(Debug, PartialEq, Clone)]
 544pub enum ContentBlock {
 545    Empty,
 546    Markdown { markdown: Entity<Markdown> },
 547    ResourceLink { resource_link: acp::ResourceLink },
 548    Image { image: Arc<gpui::Image> },
 549}
 550
 551impl ContentBlock {
 552    pub fn new(
 553        block: acp::ContentBlock,
 554        language_registry: &Arc<LanguageRegistry>,
 555        path_style: PathStyle,
 556        cx: &mut App,
 557    ) -> Self {
 558        let mut this = Self::Empty;
 559        this.append(block, language_registry, path_style, cx);
 560        this
 561    }
 562
 563    pub fn new_combined(
 564        blocks: impl IntoIterator<Item = acp::ContentBlock>,
 565        language_registry: Arc<LanguageRegistry>,
 566        path_style: PathStyle,
 567        cx: &mut App,
 568    ) -> Self {
 569        let mut this = Self::Empty;
 570        for block in blocks {
 571            this.append(block, &language_registry, path_style, cx);
 572        }
 573        this
 574    }
 575
 576    pub fn append(
 577        &mut self,
 578        block: acp::ContentBlock,
 579        language_registry: &Arc<LanguageRegistry>,
 580        path_style: PathStyle,
 581        cx: &mut App,
 582    ) {
 583        match (&mut *self, &block) {
 584            (ContentBlock::Empty, acp::ContentBlock::ResourceLink(resource_link)) => {
 585                *self = ContentBlock::ResourceLink {
 586                    resource_link: resource_link.clone(),
 587                };
 588            }
 589            (ContentBlock::Empty, acp::ContentBlock::Image(image_content)) => {
 590                if let Some(image) = Self::decode_image(image_content) {
 591                    *self = ContentBlock::Image { image };
 592                } else {
 593                    let new_content = Self::image_md(image_content);
 594                    *self = Self::create_markdown_block(new_content, language_registry, cx);
 595                }
 596            }
 597            (ContentBlock::Empty, _) => {
 598                let new_content = Self::block_string_contents(&block, path_style);
 599                *self = Self::create_markdown_block(new_content, language_registry, cx);
 600            }
 601            (ContentBlock::Markdown { markdown }, _) => {
 602                let new_content = Self::block_string_contents(&block, path_style);
 603                markdown.update(cx, |markdown, cx| markdown.append(&new_content, cx));
 604            }
 605            (ContentBlock::ResourceLink { resource_link }, _) => {
 606                let existing_content = Self::resource_link_md(&resource_link.uri, path_style);
 607                let new_content = Self::block_string_contents(&block, path_style);
 608                let combined = format!("{}\n{}", existing_content, new_content);
 609                *self = Self::create_markdown_block(combined, language_registry, cx);
 610            }
 611            (ContentBlock::Image { .. }, _) => {
 612                let new_content = Self::block_string_contents(&block, path_style);
 613                let combined = format!("`Image`\n{}", new_content);
 614                *self = Self::create_markdown_block(combined, language_registry, cx);
 615            }
 616        }
 617    }
 618
 619    fn decode_image(image_content: &acp::ImageContent) -> Option<Arc<gpui::Image>> {
 620        use base64::Engine as _;
 621
 622        let bytes = base64::engine::general_purpose::STANDARD
 623            .decode(image_content.data.as_bytes())
 624            .ok()?;
 625        let format = gpui::ImageFormat::from_mime_type(&image_content.mime_type)?;
 626        Some(Arc::new(gpui::Image::from_bytes(format, bytes)))
 627    }
 628
 629    fn create_markdown_block(
 630        content: String,
 631        language_registry: &Arc<LanguageRegistry>,
 632        cx: &mut App,
 633    ) -> ContentBlock {
 634        ContentBlock::Markdown {
 635            markdown: cx
 636                .new(|cx| Markdown::new(content.into(), Some(language_registry.clone()), None, cx)),
 637        }
 638    }
 639
 640    fn block_string_contents(block: &acp::ContentBlock, path_style: PathStyle) -> String {
 641        match block {
 642            acp::ContentBlock::Text(text_content) => text_content.text.clone(),
 643            acp::ContentBlock::ResourceLink(resource_link) => {
 644                Self::resource_link_md(&resource_link.uri, path_style)
 645            }
 646            acp::ContentBlock::Resource(acp::EmbeddedResource {
 647                resource:
 648                    acp::EmbeddedResourceResource::TextResourceContents(acp::TextResourceContents {
 649                        uri,
 650                        ..
 651                    }),
 652                ..
 653            }) => Self::resource_link_md(uri, path_style),
 654            acp::ContentBlock::Image(image) => Self::image_md(image),
 655            _ => String::new(),
 656        }
 657    }
 658
 659    fn resource_link_md(uri: &str, path_style: PathStyle) -> String {
 660        if let Some(uri) = MentionUri::parse(uri, path_style).log_err() {
 661            uri.as_link().to_string()
 662        } else {
 663            uri.to_string()
 664        }
 665    }
 666
 667    fn image_md(_image: &acp::ImageContent) -> String {
 668        "`Image`".into()
 669    }
 670
 671    pub fn to_markdown<'a>(&'a self, cx: &'a App) -> &'a str {
 672        match self {
 673            ContentBlock::Empty => "",
 674            ContentBlock::Markdown { markdown } => markdown.read(cx).source(),
 675            ContentBlock::ResourceLink { resource_link } => &resource_link.uri,
 676            ContentBlock::Image { .. } => "`Image`",
 677        }
 678    }
 679
 680    pub fn markdown(&self) -> Option<&Entity<Markdown>> {
 681        match self {
 682            ContentBlock::Empty => None,
 683            ContentBlock::Markdown { markdown } => Some(markdown),
 684            ContentBlock::ResourceLink { .. } => None,
 685            ContentBlock::Image { .. } => None,
 686        }
 687    }
 688
 689    pub fn resource_link(&self) -> Option<&acp::ResourceLink> {
 690        match self {
 691            ContentBlock::ResourceLink { resource_link } => Some(resource_link),
 692            _ => None,
 693        }
 694    }
 695
 696    pub fn image(&self) -> Option<&Arc<gpui::Image>> {
 697        match self {
 698            ContentBlock::Image { image } => Some(image),
 699            _ => None,
 700        }
 701    }
 702}
 703
 704#[derive(Debug)]
 705pub enum ToolCallContent {
 706    ContentBlock(ContentBlock),
 707    Diff(Entity<Diff>),
 708    Terminal(Entity<Terminal>),
 709}
 710
 711impl ToolCallContent {
 712    pub fn from_acp(
 713        content: acp::ToolCallContent,
 714        language_registry: Arc<LanguageRegistry>,
 715        path_style: PathStyle,
 716        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 717        cx: &mut App,
 718    ) -> Result<Option<Self>> {
 719        match content {
 720            acp::ToolCallContent::Content(acp::Content { content, .. }) => {
 721                Ok(Some(Self::ContentBlock(ContentBlock::new(
 722                    content,
 723                    &language_registry,
 724                    path_style,
 725                    cx,
 726                ))))
 727            }
 728            acp::ToolCallContent::Diff(diff) => Ok(Some(Self::Diff(cx.new(|cx| {
 729                Diff::finalized(
 730                    diff.path.to_string_lossy().into_owned(),
 731                    diff.old_text,
 732                    diff.new_text,
 733                    language_registry,
 734                    cx,
 735                )
 736            })))),
 737            acp::ToolCallContent::Terminal(acp::Terminal { terminal_id, .. }) => terminals
 738                .get(&terminal_id)
 739                .cloned()
 740                .map(|terminal| Some(Self::Terminal(terminal)))
 741                .ok_or_else(|| anyhow::anyhow!("Terminal with id `{}` not found", terminal_id)),
 742            _ => Ok(None),
 743        }
 744    }
 745
 746    pub fn update_from_acp(
 747        &mut self,
 748        new: acp::ToolCallContent,
 749        language_registry: Arc<LanguageRegistry>,
 750        path_style: PathStyle,
 751        terminals: &HashMap<acp::TerminalId, Entity<Terminal>>,
 752        cx: &mut App,
 753    ) -> Result<bool> {
 754        let needs_update = match (&self, &new) {
 755            (Self::Diff(old_diff), acp::ToolCallContent::Diff(new_diff)) => {
 756                old_diff.read(cx).needs_update(
 757                    new_diff.old_text.as_deref().unwrap_or(""),
 758                    &new_diff.new_text,
 759                    cx,
 760                )
 761            }
 762            _ => true,
 763        };
 764
 765        if let Some(update) = Self::from_acp(new, language_registry, path_style, terminals, cx)? {
 766            if needs_update {
 767                *self = update;
 768            }
 769            Ok(true)
 770        } else {
 771            Ok(false)
 772        }
 773    }
 774
 775    pub fn to_markdown(&self, cx: &App) -> String {
 776        match self {
 777            Self::ContentBlock(content) => content.to_markdown(cx).to_string(),
 778            Self::Diff(diff) => diff.read(cx).to_markdown(cx),
 779            Self::Terminal(terminal) => terminal.read(cx).to_markdown(cx),
 780        }
 781    }
 782
 783    pub fn image(&self) -> Option<&Arc<gpui::Image>> {
 784        match self {
 785            Self::ContentBlock(content) => content.image(),
 786            _ => None,
 787        }
 788    }
 789}
 790
 791#[derive(Debug, PartialEq)]
 792pub enum ToolCallUpdate {
 793    UpdateFields(acp::ToolCallUpdate),
 794    UpdateDiff(ToolCallUpdateDiff),
 795    UpdateTerminal(ToolCallUpdateTerminal),
 796}
 797
 798impl ToolCallUpdate {
 799    fn id(&self) -> &acp::ToolCallId {
 800        match self {
 801            Self::UpdateFields(update) => &update.tool_call_id,
 802            Self::UpdateDiff(diff) => &diff.id,
 803            Self::UpdateTerminal(terminal) => &terminal.id,
 804        }
 805    }
 806}
 807
 808impl From<acp::ToolCallUpdate> for ToolCallUpdate {
 809    fn from(update: acp::ToolCallUpdate) -> Self {
 810        Self::UpdateFields(update)
 811    }
 812}
 813
 814impl From<ToolCallUpdateDiff> for ToolCallUpdate {
 815    fn from(diff: ToolCallUpdateDiff) -> Self {
 816        Self::UpdateDiff(diff)
 817    }
 818}
 819
 820#[derive(Debug, PartialEq)]
 821pub struct ToolCallUpdateDiff {
 822    pub id: acp::ToolCallId,
 823    pub diff: Entity<Diff>,
 824}
 825
 826impl From<ToolCallUpdateTerminal> for ToolCallUpdate {
 827    fn from(terminal: ToolCallUpdateTerminal) -> Self {
 828        Self::UpdateTerminal(terminal)
 829    }
 830}
 831
 832#[derive(Debug, PartialEq)]
 833pub struct ToolCallUpdateTerminal {
 834    pub id: acp::ToolCallId,
 835    pub terminal: Entity<Terminal>,
 836}
 837
 838#[derive(Debug, Default)]
 839pub struct Plan {
 840    pub entries: Vec<PlanEntry>,
 841}
 842
 843#[derive(Debug)]
 844pub struct PlanStats<'a> {
 845    pub in_progress_entry: Option<&'a PlanEntry>,
 846    pub pending: u32,
 847    pub completed: u32,
 848}
 849
 850impl Plan {
 851    pub fn is_empty(&self) -> bool {
 852        self.entries.is_empty()
 853    }
 854
 855    pub fn stats(&self) -> PlanStats<'_> {
 856        let mut stats = PlanStats {
 857            in_progress_entry: None,
 858            pending: 0,
 859            completed: 0,
 860        };
 861
 862        for entry in &self.entries {
 863            match &entry.status {
 864                acp::PlanEntryStatus::Pending => {
 865                    stats.pending += 1;
 866                }
 867                acp::PlanEntryStatus::InProgress => {
 868                    stats.in_progress_entry = stats.in_progress_entry.or(Some(entry));
 869                }
 870                acp::PlanEntryStatus::Completed => {
 871                    stats.completed += 1;
 872                }
 873                _ => {}
 874            }
 875        }
 876
 877        stats
 878    }
 879}
 880
 881#[derive(Debug)]
 882pub struct PlanEntry {
 883    pub content: Entity<Markdown>,
 884    pub priority: acp::PlanEntryPriority,
 885    pub status: acp::PlanEntryStatus,
 886}
 887
 888impl PlanEntry {
 889    pub fn from_acp(entry: acp::PlanEntry, cx: &mut App) -> Self {
 890        Self {
 891            content: cx.new(|cx| Markdown::new(entry.content.into(), None, None, cx)),
 892            priority: entry.priority,
 893            status: entry.status,
 894        }
 895    }
 896}
 897
 898#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 899pub struct TokenUsage {
 900    pub max_tokens: u64,
 901    pub used_tokens: u64,
 902    pub input_tokens: u64,
 903    pub output_tokens: u64,
 904    pub max_output_tokens: Option<u64>,
 905}
 906
 907pub const TOKEN_USAGE_WARNING_THRESHOLD: f32 = 0.8;
 908
 909impl TokenUsage {
 910    pub fn ratio(&self) -> TokenUsageRatio {
 911        #[cfg(debug_assertions)]
 912        let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD")
 913            .unwrap_or(TOKEN_USAGE_WARNING_THRESHOLD.to_string())
 914            .parse()
 915            .unwrap();
 916        #[cfg(not(debug_assertions))]
 917        let warning_threshold: f32 = TOKEN_USAGE_WARNING_THRESHOLD;
 918
 919        // When the maximum is unknown because there is no selected model,
 920        // avoid showing the token limit warning.
 921        if self.max_tokens == 0 {
 922            TokenUsageRatio::Normal
 923        } else if self.used_tokens >= self.max_tokens {
 924            TokenUsageRatio::Exceeded
 925        } else if self.used_tokens as f32 / self.max_tokens as f32 >= warning_threshold {
 926            TokenUsageRatio::Warning
 927        } else {
 928            TokenUsageRatio::Normal
 929        }
 930    }
 931}
 932
 933#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 934pub enum TokenUsageRatio {
 935    Normal,
 936    Warning,
 937    Exceeded,
 938}
 939
 940#[derive(Debug, Clone)]
 941pub struct RetryStatus {
 942    pub last_error: SharedString,
 943    pub attempt: usize,
 944    pub max_attempts: usize,
 945    pub started_at: Instant,
 946    pub duration: Duration,
 947}
 948
 949struct RunningTurn {
 950    id: u32,
 951    send_task: Task<()>,
 952}
 953
 954pub struct AcpThread {
 955    parent_session_id: Option<acp::SessionId>,
 956    title: SharedString,
 957    entries: Vec<AgentThreadEntry>,
 958    plan: Plan,
 959    project: Entity<Project>,
 960    action_log: Entity<ActionLog>,
 961    shared_buffers: HashMap<Entity<Buffer>, BufferSnapshot>,
 962    turn_id: u32,
 963    running_turn: Option<RunningTurn>,
 964    connection: Rc<dyn AgentConnection>,
 965    session_id: acp::SessionId,
 966    token_usage: Option<TokenUsage>,
 967    prompt_capabilities: acp::PromptCapabilities,
 968    _observe_prompt_capabilities: Task<anyhow::Result<()>>,
 969    terminals: HashMap<acp::TerminalId, Entity<Terminal>>,
 970    pending_terminal_output: HashMap<acp::TerminalId, Vec<Vec<u8>>>,
 971    pending_terminal_exit: HashMap<acp::TerminalId, acp::TerminalExitStatus>,
 972    had_error: bool,
 973}
 974
 975impl From<&AcpThread> for ActionLogTelemetry {
 976    fn from(value: &AcpThread) -> Self {
 977        Self {
 978            agent_telemetry_id: value.connection().telemetry_id(),
 979            session_id: value.session_id.0.clone(),
 980        }
 981    }
 982}
 983
 984#[derive(Debug)]
 985pub enum AcpThreadEvent {
 986    NewEntry,
 987    TitleUpdated,
 988    TokenUsageUpdated,
 989    EntryUpdated(usize),
 990    EntriesRemoved(Range<usize>),
 991    ToolAuthorizationRequested(acp::ToolCallId),
 992    ToolAuthorizationReceived(acp::ToolCallId),
 993    Retry(RetryStatus),
 994    SubagentSpawned(acp::SessionId),
 995    Stopped(acp::StopReason),
 996    Error,
 997    LoadError(LoadError),
 998    PromptCapabilitiesUpdated,
 999    Refusal,
1000    AvailableCommandsUpdated(Vec<acp::AvailableCommand>),
1001    ModeUpdated(acp::SessionModeId),
1002    ConfigOptionsUpdated(Vec<acp::SessionConfigOption>),
1003}
1004
1005impl EventEmitter<AcpThreadEvent> for AcpThread {}
1006
1007#[derive(Debug, Clone)]
1008pub enum TerminalProviderEvent {
1009    Created {
1010        terminal_id: acp::TerminalId,
1011        label: String,
1012        cwd: Option<PathBuf>,
1013        output_byte_limit: Option<u64>,
1014        terminal: Entity<::terminal::Terminal>,
1015    },
1016    Output {
1017        terminal_id: acp::TerminalId,
1018        data: Vec<u8>,
1019    },
1020    TitleChanged {
1021        terminal_id: acp::TerminalId,
1022        title: String,
1023    },
1024    Exit {
1025        terminal_id: acp::TerminalId,
1026        status: acp::TerminalExitStatus,
1027    },
1028}
1029
1030#[derive(Debug, Clone)]
1031pub enum TerminalProviderCommand {
1032    WriteInput {
1033        terminal_id: acp::TerminalId,
1034        bytes: Vec<u8>,
1035    },
1036    Resize {
1037        terminal_id: acp::TerminalId,
1038        cols: u16,
1039        rows: u16,
1040    },
1041    Close {
1042        terminal_id: acp::TerminalId,
1043    },
1044}
1045
1046impl AcpThread {
1047    pub fn on_terminal_provider_event(
1048        &mut self,
1049        event: TerminalProviderEvent,
1050        cx: &mut Context<Self>,
1051    ) {
1052        match event {
1053            TerminalProviderEvent::Created {
1054                terminal_id,
1055                label,
1056                cwd,
1057                output_byte_limit,
1058                terminal,
1059            } => {
1060                let entity = self.register_terminal_created(
1061                    terminal_id.clone(),
1062                    label,
1063                    cwd,
1064                    output_byte_limit,
1065                    terminal,
1066                    cx,
1067                );
1068
1069                if let Some(mut chunks) = self.pending_terminal_output.remove(&terminal_id) {
1070                    for data in chunks.drain(..) {
1071                        entity.update(cx, |term, cx| {
1072                            term.inner().update(cx, |inner, cx| {
1073                                inner.write_output(&data, cx);
1074                            })
1075                        });
1076                    }
1077                }
1078
1079                if let Some(_status) = self.pending_terminal_exit.remove(&terminal_id) {
1080                    entity.update(cx, |_term, cx| {
1081                        cx.notify();
1082                    });
1083                }
1084
1085                cx.notify();
1086            }
1087            TerminalProviderEvent::Output { terminal_id, data } => {
1088                if let Some(entity) = self.terminals.get(&terminal_id) {
1089                    entity.update(cx, |term, cx| {
1090                        term.inner().update(cx, |inner, cx| {
1091                            inner.write_output(&data, cx);
1092                        })
1093                    });
1094                } else {
1095                    self.pending_terminal_output
1096                        .entry(terminal_id)
1097                        .or_default()
1098                        .push(data);
1099                }
1100            }
1101            TerminalProviderEvent::TitleChanged { terminal_id, title } => {
1102                if let Some(entity) = self.terminals.get(&terminal_id) {
1103                    entity.update(cx, |term, cx| {
1104                        term.inner().update(cx, |inner, cx| {
1105                            inner.breadcrumb_text = title;
1106                            cx.emit(::terminal::Event::BreadcrumbsChanged);
1107                        })
1108                    });
1109                }
1110            }
1111            TerminalProviderEvent::Exit {
1112                terminal_id,
1113                status,
1114            } => {
1115                if let Some(entity) = self.terminals.get(&terminal_id) {
1116                    entity.update(cx, |_term, cx| {
1117                        cx.notify();
1118                    });
1119                } else {
1120                    self.pending_terminal_exit.insert(terminal_id, status);
1121                }
1122            }
1123        }
1124    }
1125}
1126
1127#[derive(PartialEq, Eq, Debug)]
1128pub enum ThreadStatus {
1129    Idle,
1130    Generating,
1131}
1132
1133#[derive(Debug, Clone)]
1134pub enum LoadError {
1135    Unsupported {
1136        command: SharedString,
1137        current_version: SharedString,
1138        minimum_version: SharedString,
1139    },
1140    FailedToInstall(SharedString),
1141    Exited {
1142        status: ExitStatus,
1143    },
1144    Other(SharedString),
1145}
1146
1147impl Display for LoadError {
1148    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1149        match self {
1150            LoadError::Unsupported {
1151                command: path,
1152                current_version,
1153                minimum_version,
1154            } => {
1155                write!(
1156                    f,
1157                    "version {current_version} from {path} is not supported (need at least {minimum_version})"
1158                )
1159            }
1160            LoadError::FailedToInstall(msg) => write!(f, "Failed to install: {msg}"),
1161            LoadError::Exited { status } => write!(f, "Server exited with status {status}"),
1162            LoadError::Other(msg) => write!(f, "{msg}"),
1163        }
1164    }
1165}
1166
1167impl Error for LoadError {}
1168
1169impl AcpThread {
1170    pub fn new(
1171        parent_session_id: Option<acp::SessionId>,
1172        title: impl Into<SharedString>,
1173        connection: Rc<dyn AgentConnection>,
1174        project: Entity<Project>,
1175        action_log: Entity<ActionLog>,
1176        session_id: acp::SessionId,
1177        mut prompt_capabilities_rx: watch::Receiver<acp::PromptCapabilities>,
1178        cx: &mut Context<Self>,
1179    ) -> Self {
1180        let prompt_capabilities = prompt_capabilities_rx.borrow().clone();
1181        let task = cx.spawn::<_, anyhow::Result<()>>(async move |this, cx| {
1182            loop {
1183                let caps = prompt_capabilities_rx.recv().await?;
1184                this.update(cx, |this, cx| {
1185                    this.prompt_capabilities = caps;
1186                    cx.emit(AcpThreadEvent::PromptCapabilitiesUpdated);
1187                })?;
1188            }
1189        });
1190
1191        Self {
1192            parent_session_id,
1193            action_log,
1194            shared_buffers: Default::default(),
1195            entries: Default::default(),
1196            plan: Default::default(),
1197            title: title.into(),
1198            project,
1199            running_turn: None,
1200            turn_id: 0,
1201            connection,
1202            session_id,
1203            token_usage: None,
1204            prompt_capabilities,
1205            _observe_prompt_capabilities: task,
1206            terminals: HashMap::default(),
1207            pending_terminal_output: HashMap::default(),
1208            pending_terminal_exit: HashMap::default(),
1209            had_error: false,
1210        }
1211    }
1212
1213    pub fn parent_session_id(&self) -> Option<&acp::SessionId> {
1214        self.parent_session_id.as_ref()
1215    }
1216
1217    pub fn prompt_capabilities(&self) -> acp::PromptCapabilities {
1218        self.prompt_capabilities.clone()
1219    }
1220
1221    pub fn connection(&self) -> &Rc<dyn AgentConnection> {
1222        &self.connection
1223    }
1224
1225    pub fn action_log(&self) -> &Entity<ActionLog> {
1226        &self.action_log
1227    }
1228
1229    pub fn project(&self) -> &Entity<Project> {
1230        &self.project
1231    }
1232
1233    pub fn title(&self) -> SharedString {
1234        self.title.clone()
1235    }
1236
1237    pub fn entries(&self) -> &[AgentThreadEntry] {
1238        &self.entries
1239    }
1240
1241    pub fn session_id(&self) -> &acp::SessionId {
1242        &self.session_id
1243    }
1244
1245    pub fn status(&self) -> ThreadStatus {
1246        if self.running_turn.is_some() {
1247            ThreadStatus::Generating
1248        } else {
1249            ThreadStatus::Idle
1250        }
1251    }
1252
1253    pub fn had_error(&self) -> bool {
1254        self.had_error
1255    }
1256
1257    pub fn is_waiting_for_confirmation(&self) -> bool {
1258        for entry in self.entries.iter().rev() {
1259            match entry {
1260                AgentThreadEntry::UserMessage(_) => return false,
1261                AgentThreadEntry::ToolCall(ToolCall {
1262                    status: ToolCallStatus::WaitingForConfirmation { .. },
1263                    ..
1264                }) => return true,
1265                AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
1266            }
1267        }
1268        false
1269    }
1270
1271    pub fn token_usage(&self) -> Option<&TokenUsage> {
1272        self.token_usage.as_ref()
1273    }
1274
1275    pub fn has_pending_edit_tool_calls(&self) -> bool {
1276        for entry in self.entries.iter().rev() {
1277            match entry {
1278                AgentThreadEntry::UserMessage(_) => return false,
1279                AgentThreadEntry::ToolCall(
1280                    call @ ToolCall {
1281                        status: ToolCallStatus::InProgress | ToolCallStatus::Pending,
1282                        ..
1283                    },
1284                ) if call.diffs().next().is_some() => {
1285                    return true;
1286                }
1287                AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
1288            }
1289        }
1290
1291        false
1292    }
1293
1294    pub fn has_in_progress_tool_calls(&self) -> bool {
1295        for entry in self.entries.iter().rev() {
1296            match entry {
1297                AgentThreadEntry::UserMessage(_) => return false,
1298                AgentThreadEntry::ToolCall(ToolCall {
1299                    status: ToolCallStatus::InProgress | ToolCallStatus::Pending,
1300                    ..
1301                }) => {
1302                    return true;
1303                }
1304                AgentThreadEntry::ToolCall(_) | AgentThreadEntry::AssistantMessage(_) => {}
1305            }
1306        }
1307
1308        false
1309    }
1310
1311    pub fn used_tools_since_last_user_message(&self) -> bool {
1312        for entry in self.entries.iter().rev() {
1313            match entry {
1314                AgentThreadEntry::UserMessage(..) => return false,
1315                AgentThreadEntry::AssistantMessage(..) => continue,
1316                AgentThreadEntry::ToolCall(..) => return true,
1317            }
1318        }
1319
1320        false
1321    }
1322
1323    pub fn handle_session_update(
1324        &mut self,
1325        update: acp::SessionUpdate,
1326        cx: &mut Context<Self>,
1327    ) -> Result<(), acp::Error> {
1328        match update {
1329            acp::SessionUpdate::UserMessageChunk(acp::ContentChunk { content, .. }) => {
1330                self.push_user_content_block(None, content, cx);
1331            }
1332            acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk { content, .. }) => {
1333                self.push_assistant_content_block(content, false, cx);
1334            }
1335            acp::SessionUpdate::AgentThoughtChunk(acp::ContentChunk { content, .. }) => {
1336                self.push_assistant_content_block(content, true, cx);
1337            }
1338            acp::SessionUpdate::ToolCall(tool_call) => {
1339                self.upsert_tool_call(tool_call, cx)?;
1340            }
1341            acp::SessionUpdate::ToolCallUpdate(tool_call_update) => {
1342                self.update_tool_call(tool_call_update, cx)?;
1343            }
1344            acp::SessionUpdate::Plan(plan) => {
1345                self.update_plan(plan, cx);
1346            }
1347            acp::SessionUpdate::AvailableCommandsUpdate(acp::AvailableCommandsUpdate {
1348                available_commands,
1349                ..
1350            }) => cx.emit(AcpThreadEvent::AvailableCommandsUpdated(available_commands)),
1351            acp::SessionUpdate::CurrentModeUpdate(acp::CurrentModeUpdate {
1352                current_mode_id,
1353                ..
1354            }) => cx.emit(AcpThreadEvent::ModeUpdated(current_mode_id)),
1355            acp::SessionUpdate::ConfigOptionUpdate(acp::ConfigOptionUpdate {
1356                config_options,
1357                ..
1358            }) => cx.emit(AcpThreadEvent::ConfigOptionsUpdated(config_options)),
1359            _ => {}
1360        }
1361        Ok(())
1362    }
1363
1364    pub fn push_user_content_block(
1365        &mut self,
1366        message_id: Option<UserMessageId>,
1367        chunk: acp::ContentBlock,
1368        cx: &mut Context<Self>,
1369    ) {
1370        self.push_user_content_block_with_indent(message_id, chunk, false, cx)
1371    }
1372
1373    pub fn push_user_content_block_with_indent(
1374        &mut self,
1375        message_id: Option<UserMessageId>,
1376        chunk: acp::ContentBlock,
1377        indented: bool,
1378        cx: &mut Context<Self>,
1379    ) {
1380        let language_registry = self.project.read(cx).languages().clone();
1381        let path_style = self.project.read(cx).path_style(cx);
1382        let entries_len = self.entries.len();
1383
1384        if let Some(last_entry) = self.entries.last_mut()
1385            && let AgentThreadEntry::UserMessage(UserMessage {
1386                id,
1387                content,
1388                chunks,
1389                indented: existing_indented,
1390                ..
1391            }) = last_entry
1392            && *existing_indented == indented
1393        {
1394            *id = message_id.or(id.take());
1395            content.append(chunk.clone(), &language_registry, path_style, cx);
1396            chunks.push(chunk);
1397            let idx = entries_len - 1;
1398            cx.emit(AcpThreadEvent::EntryUpdated(idx));
1399        } else {
1400            let content = ContentBlock::new(chunk.clone(), &language_registry, path_style, cx);
1401            self.push_entry(
1402                AgentThreadEntry::UserMessage(UserMessage {
1403                    id: message_id,
1404                    content,
1405                    chunks: vec![chunk],
1406                    checkpoint: None,
1407                    indented,
1408                }),
1409                cx,
1410            );
1411        }
1412    }
1413
1414    pub fn push_assistant_content_block(
1415        &mut self,
1416        chunk: acp::ContentBlock,
1417        is_thought: bool,
1418        cx: &mut Context<Self>,
1419    ) {
1420        self.push_assistant_content_block_with_indent(chunk, is_thought, false, cx)
1421    }
1422
1423    pub fn push_assistant_content_block_with_indent(
1424        &mut self,
1425        chunk: acp::ContentBlock,
1426        is_thought: bool,
1427        indented: bool,
1428        cx: &mut Context<Self>,
1429    ) {
1430        let language_registry = self.project.read(cx).languages().clone();
1431        let path_style = self.project.read(cx).path_style(cx);
1432        let entries_len = self.entries.len();
1433        if let Some(last_entry) = self.entries.last_mut()
1434            && let AgentThreadEntry::AssistantMessage(AssistantMessage {
1435                chunks,
1436                indented: existing_indented,
1437                is_subagent_output: _,
1438            }) = last_entry
1439            && *existing_indented == indented
1440        {
1441            let idx = entries_len - 1;
1442            cx.emit(AcpThreadEvent::EntryUpdated(idx));
1443            match (chunks.last_mut(), is_thought) {
1444                (Some(AssistantMessageChunk::Message { block }), false)
1445                | (Some(AssistantMessageChunk::Thought { block }), true) => {
1446                    block.append(chunk, &language_registry, path_style, cx)
1447                }
1448                _ => {
1449                    let block = ContentBlock::new(chunk, &language_registry, path_style, cx);
1450                    if is_thought {
1451                        chunks.push(AssistantMessageChunk::Thought { block })
1452                    } else {
1453                        chunks.push(AssistantMessageChunk::Message { block })
1454                    }
1455                }
1456            }
1457        } else {
1458            let block = ContentBlock::new(chunk, &language_registry, path_style, cx);
1459            let chunk = if is_thought {
1460                AssistantMessageChunk::Thought { block }
1461            } else {
1462                AssistantMessageChunk::Message { block }
1463            };
1464
1465            self.push_entry(
1466                AgentThreadEntry::AssistantMessage(AssistantMessage {
1467                    chunks: vec![chunk],
1468                    indented,
1469                    is_subagent_output: false,
1470                }),
1471                cx,
1472            );
1473        }
1474    }
1475
1476    fn push_entry(&mut self, entry: AgentThreadEntry, cx: &mut Context<Self>) {
1477        self.entries.push(entry);
1478        cx.emit(AcpThreadEvent::NewEntry);
1479    }
1480
1481    pub fn can_set_title(&mut self, cx: &mut Context<Self>) -> bool {
1482        self.connection.set_title(&self.session_id, cx).is_some()
1483    }
1484
1485    pub fn set_title(&mut self, title: SharedString, cx: &mut Context<Self>) -> Task<Result<()>> {
1486        if title != self.title {
1487            self.title = title.clone();
1488            cx.emit(AcpThreadEvent::TitleUpdated);
1489            if let Some(set_title) = self.connection.set_title(&self.session_id, cx) {
1490                return set_title.run(title, cx);
1491            }
1492        }
1493        Task::ready(Ok(()))
1494    }
1495
1496    pub fn subagent_spawned(&mut self, session_id: acp::SessionId, cx: &mut Context<Self>) {
1497        cx.emit(AcpThreadEvent::SubagentSpawned(session_id));
1498    }
1499
1500    pub fn update_token_usage(&mut self, usage: Option<TokenUsage>, cx: &mut Context<Self>) {
1501        self.token_usage = usage;
1502        cx.emit(AcpThreadEvent::TokenUsageUpdated);
1503    }
1504
1505    pub fn update_retry_status(&mut self, status: RetryStatus, cx: &mut Context<Self>) {
1506        cx.emit(AcpThreadEvent::Retry(status));
1507    }
1508
1509    pub fn update_tool_call(
1510        &mut self,
1511        update: impl Into<ToolCallUpdate>,
1512        cx: &mut Context<Self>,
1513    ) -> Result<()> {
1514        let update = update.into();
1515        let languages = self.project.read(cx).languages().clone();
1516        let path_style = self.project.read(cx).path_style(cx);
1517
1518        let ix = match self.index_for_tool_call(update.id()) {
1519            Some(ix) => ix,
1520            None => {
1521                // Tool call not found - create a failed tool call entry
1522                let failed_tool_call = ToolCall {
1523                    id: update.id().clone(),
1524                    label: cx.new(|cx| Markdown::new("Tool call not found".into(), None, None, cx)),
1525                    kind: acp::ToolKind::Fetch,
1526                    content: vec![ToolCallContent::ContentBlock(ContentBlock::new(
1527                        "Tool call not found".into(),
1528                        &languages,
1529                        path_style,
1530                        cx,
1531                    ))],
1532                    status: ToolCallStatus::Failed,
1533                    locations: Vec::new(),
1534                    resolved_locations: Vec::new(),
1535                    raw_input: None,
1536                    raw_input_markdown: None,
1537                    raw_output: None,
1538                    tool_name: None,
1539                    subagent_session_info: None,
1540                };
1541                self.push_entry(AgentThreadEntry::ToolCall(failed_tool_call), cx);
1542                return Ok(());
1543            }
1544        };
1545        let AgentThreadEntry::ToolCall(call) = &mut self.entries[ix] else {
1546            unreachable!()
1547        };
1548
1549        match update {
1550            ToolCallUpdate::UpdateFields(update) => {
1551                let location_updated = update.fields.locations.is_some();
1552                call.update_fields(
1553                    update.fields,
1554                    update.meta,
1555                    languages,
1556                    path_style,
1557                    &self.terminals,
1558                    cx,
1559                )?;
1560                if location_updated {
1561                    self.resolve_locations(update.tool_call_id, cx);
1562                }
1563            }
1564            ToolCallUpdate::UpdateDiff(update) => {
1565                call.content.clear();
1566                call.content.push(ToolCallContent::Diff(update.diff));
1567            }
1568            ToolCallUpdate::UpdateTerminal(update) => {
1569                call.content.clear();
1570                call.content
1571                    .push(ToolCallContent::Terminal(update.terminal));
1572            }
1573        }
1574
1575        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1576
1577        Ok(())
1578    }
1579
1580    /// Updates a tool call if id matches an existing entry, otherwise inserts a new one.
1581    pub fn upsert_tool_call(
1582        &mut self,
1583        tool_call: acp::ToolCall,
1584        cx: &mut Context<Self>,
1585    ) -> Result<(), acp::Error> {
1586        let status = tool_call.status.into();
1587        self.upsert_tool_call_inner(tool_call.into(), status, cx)
1588    }
1589
1590    /// Fails if id does not match an existing entry.
1591    pub fn upsert_tool_call_inner(
1592        &mut self,
1593        update: acp::ToolCallUpdate,
1594        status: ToolCallStatus,
1595        cx: &mut Context<Self>,
1596    ) -> Result<(), acp::Error> {
1597        let language_registry = self.project.read(cx).languages().clone();
1598        let path_style = self.project.read(cx).path_style(cx);
1599        let id = update.tool_call_id.clone();
1600
1601        let agent_telemetry_id = self.connection().telemetry_id();
1602        let session = self.session_id();
1603        if let ToolCallStatus::Completed | ToolCallStatus::Failed = status {
1604            let status = if matches!(status, ToolCallStatus::Completed) {
1605                "completed"
1606            } else {
1607                "failed"
1608            };
1609            telemetry::event!(
1610                "Agent Tool Call Completed",
1611                agent_telemetry_id,
1612                session,
1613                status
1614            );
1615        }
1616
1617        if let Some(ix) = self.index_for_tool_call(&id) {
1618            let AgentThreadEntry::ToolCall(call) = &mut self.entries[ix] else {
1619                unreachable!()
1620            };
1621
1622            call.update_fields(
1623                update.fields,
1624                update.meta,
1625                language_registry,
1626                path_style,
1627                &self.terminals,
1628                cx,
1629            )?;
1630            call.status = status;
1631
1632            cx.emit(AcpThreadEvent::EntryUpdated(ix));
1633        } else {
1634            let call = ToolCall::from_acp(
1635                update.try_into()?,
1636                status,
1637                language_registry,
1638                self.project.read(cx).path_style(cx),
1639                &self.terminals,
1640                cx,
1641            )?;
1642            self.push_entry(AgentThreadEntry::ToolCall(call), cx);
1643        };
1644
1645        self.resolve_locations(id, cx);
1646        Ok(())
1647    }
1648
1649    fn index_for_tool_call(&self, id: &acp::ToolCallId) -> Option<usize> {
1650        self.entries
1651            .iter()
1652            .enumerate()
1653            .rev()
1654            .find_map(|(index, entry)| {
1655                if let AgentThreadEntry::ToolCall(tool_call) = entry
1656                    && &tool_call.id == id
1657                {
1658                    Some(index)
1659                } else {
1660                    None
1661                }
1662            })
1663    }
1664
1665    fn tool_call_mut(&mut self, id: &acp::ToolCallId) -> Option<(usize, &mut ToolCall)> {
1666        // The tool call we are looking for is typically the last one, or very close to the end.
1667        // At the moment, it doesn't seem like a hashmap would be a good fit for this use case.
1668        self.entries
1669            .iter_mut()
1670            .enumerate()
1671            .rev()
1672            .find_map(|(index, tool_call)| {
1673                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1674                    && &tool_call.id == id
1675                {
1676                    Some((index, tool_call))
1677                } else {
1678                    None
1679                }
1680            })
1681    }
1682
1683    pub fn tool_call(&self, id: &acp::ToolCallId) -> Option<(usize, &ToolCall)> {
1684        self.entries
1685            .iter()
1686            .enumerate()
1687            .rev()
1688            .find_map(|(index, tool_call)| {
1689                if let AgentThreadEntry::ToolCall(tool_call) = tool_call
1690                    && &tool_call.id == id
1691                {
1692                    Some((index, tool_call))
1693                } else {
1694                    None
1695                }
1696            })
1697    }
1698
1699    pub fn tool_call_for_subagent(&self, session_id: &acp::SessionId) -> Option<&ToolCall> {
1700        self.entries.iter().find_map(|entry| match entry {
1701            AgentThreadEntry::ToolCall(tool_call) => {
1702                if let Some(subagent_session_info) = &tool_call.subagent_session_info
1703                    && &subagent_session_info.session_id == session_id
1704                {
1705                    Some(tool_call)
1706                } else {
1707                    None
1708                }
1709            }
1710            _ => None,
1711        })
1712    }
1713
1714    pub fn resolve_locations(&mut self, id: acp::ToolCallId, cx: &mut Context<Self>) {
1715        let project = self.project.clone();
1716        let should_update_agent_location = self.parent_session_id.is_none();
1717        let Some((_, tool_call)) = self.tool_call_mut(&id) else {
1718            return;
1719        };
1720        let task = tool_call.resolve_locations(project, cx);
1721        cx.spawn(async move |this, cx| {
1722            let resolved_locations = task.await;
1723
1724            this.update(cx, |this, cx| {
1725                let project = this.project.clone();
1726
1727                for location in resolved_locations.iter().flatten() {
1728                    this.shared_buffers
1729                        .insert(location.buffer.clone(), location.buffer.read(cx).snapshot());
1730                }
1731                let Some((ix, tool_call)) = this.tool_call_mut(&id) else {
1732                    return;
1733                };
1734
1735                if let Some(Some(location)) = resolved_locations.last() {
1736                    project.update(cx, |project, cx| {
1737                        let should_ignore = if let Some(agent_location) = project
1738                            .agent_location()
1739                            .filter(|agent_location| agent_location.buffer == location.buffer)
1740                        {
1741                            let snapshot = location.buffer.read(cx).snapshot();
1742                            let old_position = agent_location.position.to_point(&snapshot);
1743                            let new_position = location.position.to_point(&snapshot);
1744
1745                            // ignore this so that when we get updates from the edit tool
1746                            // the position doesn't reset to the startof line
1747                            old_position.row == new_position.row
1748                                && old_position.column > new_position.column
1749                        } else {
1750                            false
1751                        };
1752                        if !should_ignore && should_update_agent_location {
1753                            project.set_agent_location(Some(location.into()), cx);
1754                        }
1755                    });
1756                }
1757
1758                let resolved_locations = resolved_locations
1759                    .iter()
1760                    .map(|l| l.as_ref().map(|l| AgentLocation::from(l)))
1761                    .collect::<Vec<_>>();
1762
1763                if tool_call.resolved_locations != resolved_locations {
1764                    tool_call.resolved_locations = resolved_locations;
1765                    cx.emit(AcpThreadEvent::EntryUpdated(ix));
1766                }
1767            })
1768        })
1769        .detach();
1770    }
1771
1772    pub fn request_tool_call_authorization(
1773        &mut self,
1774        tool_call: acp::ToolCallUpdate,
1775        options: PermissionOptions,
1776        cx: &mut Context<Self>,
1777    ) -> Result<Task<acp::RequestPermissionOutcome>> {
1778        let (tx, rx) = oneshot::channel();
1779
1780        let status = ToolCallStatus::WaitingForConfirmation {
1781            options,
1782            respond_tx: tx,
1783        };
1784
1785        let tool_call_id = tool_call.tool_call_id.clone();
1786        self.upsert_tool_call_inner(tool_call, status, cx)?;
1787        cx.emit(AcpThreadEvent::ToolAuthorizationRequested(
1788            tool_call_id.clone(),
1789        ));
1790
1791        Ok(cx.spawn(async move |this, cx| {
1792            let outcome = match rx.await {
1793                Ok(option) => acp::RequestPermissionOutcome::Selected(
1794                    acp::SelectedPermissionOutcome::new(option),
1795                ),
1796                Err(oneshot::Canceled) => acp::RequestPermissionOutcome::Cancelled,
1797            };
1798            this.update(cx, |_this, cx| {
1799                cx.emit(AcpThreadEvent::ToolAuthorizationReceived(tool_call_id))
1800            })
1801            .ok();
1802            outcome
1803        }))
1804    }
1805
1806    pub fn authorize_tool_call(
1807        &mut self,
1808        id: acp::ToolCallId,
1809        option_id: acp::PermissionOptionId,
1810        option_kind: acp::PermissionOptionKind,
1811        cx: &mut Context<Self>,
1812    ) {
1813        let Some((ix, call)) = self.tool_call_mut(&id) else {
1814            return;
1815        };
1816
1817        let new_status = match option_kind {
1818            acp::PermissionOptionKind::RejectOnce | acp::PermissionOptionKind::RejectAlways => {
1819                ToolCallStatus::Rejected
1820            }
1821            acp::PermissionOptionKind::AllowOnce | acp::PermissionOptionKind::AllowAlways => {
1822                ToolCallStatus::InProgress
1823            }
1824            _ => ToolCallStatus::InProgress,
1825        };
1826
1827        let curr_status = mem::replace(&mut call.status, new_status);
1828
1829        if let ToolCallStatus::WaitingForConfirmation { respond_tx, .. } = curr_status {
1830            respond_tx.send(option_id).log_err();
1831        } else if cfg!(debug_assertions) {
1832            panic!("tried to authorize an already authorized tool call");
1833        }
1834
1835        cx.emit(AcpThreadEvent::EntryUpdated(ix));
1836    }
1837
1838    pub fn plan(&self) -> &Plan {
1839        &self.plan
1840    }
1841
1842    pub fn update_plan(&mut self, request: acp::Plan, cx: &mut Context<Self>) {
1843        let new_entries_len = request.entries.len();
1844        let mut new_entries = request.entries.into_iter();
1845
1846        // Reuse existing markdown to prevent flickering
1847        for (old, new) in self.plan.entries.iter_mut().zip(new_entries.by_ref()) {
1848            let PlanEntry {
1849                content,
1850                priority,
1851                status,
1852            } = old;
1853            content.update(cx, |old, cx| {
1854                old.replace(new.content, cx);
1855            });
1856            *priority = new.priority;
1857            *status = new.status;
1858        }
1859        for new in new_entries {
1860            self.plan.entries.push(PlanEntry::from_acp(new, cx))
1861        }
1862        self.plan.entries.truncate(new_entries_len);
1863
1864        cx.notify();
1865    }
1866
1867    fn clear_completed_plan_entries(&mut self, cx: &mut Context<Self>) {
1868        self.plan
1869            .entries
1870            .retain(|entry| !matches!(entry.status, acp::PlanEntryStatus::Completed));
1871        cx.notify();
1872    }
1873
1874    #[cfg(any(test, feature = "test-support"))]
1875    pub fn send_raw(
1876        &mut self,
1877        message: &str,
1878        cx: &mut Context<Self>,
1879    ) -> BoxFuture<'static, Result<Option<acp::PromptResponse>>> {
1880        self.send(vec![message.into()], cx)
1881    }
1882
1883    pub fn send(
1884        &mut self,
1885        message: Vec<acp::ContentBlock>,
1886        cx: &mut Context<Self>,
1887    ) -> BoxFuture<'static, Result<Option<acp::PromptResponse>>> {
1888        let block = ContentBlock::new_combined(
1889            message.clone(),
1890            self.project.read(cx).languages().clone(),
1891            self.project.read(cx).path_style(cx),
1892            cx,
1893        );
1894        let request = acp::PromptRequest::new(self.session_id.clone(), message.clone());
1895        let git_store = self.project.read(cx).git_store().clone();
1896
1897        let message_id = if self.connection.truncate(&self.session_id, cx).is_some() {
1898            Some(UserMessageId::new())
1899        } else {
1900            None
1901        };
1902
1903        self.run_turn(cx, async move |this, cx| {
1904            this.update(cx, |this, cx| {
1905                this.push_entry(
1906                    AgentThreadEntry::UserMessage(UserMessage {
1907                        id: message_id.clone(),
1908                        content: block,
1909                        chunks: message,
1910                        checkpoint: None,
1911                        indented: false,
1912                    }),
1913                    cx,
1914                );
1915            })
1916            .ok();
1917
1918            let old_checkpoint = git_store
1919                .update(cx, |git, cx| git.checkpoint(cx))
1920                .await
1921                .context("failed to get old checkpoint")
1922                .log_err();
1923            this.update(cx, |this, cx| {
1924                if let Some((_ix, message)) = this.last_user_message() {
1925                    message.checkpoint = old_checkpoint.map(|git_checkpoint| Checkpoint {
1926                        git_checkpoint,
1927                        show: false,
1928                    });
1929                }
1930                this.connection.prompt(message_id, request, cx)
1931            })?
1932            .await
1933        })
1934    }
1935
1936    pub fn can_retry(&self, cx: &App) -> bool {
1937        self.connection.retry(&self.session_id, cx).is_some()
1938    }
1939
1940    pub fn retry(
1941        &mut self,
1942        cx: &mut Context<Self>,
1943    ) -> BoxFuture<'static, Result<Option<acp::PromptResponse>>> {
1944        self.run_turn(cx, async move |this, cx| {
1945            this.update(cx, |this, cx| {
1946                this.connection
1947                    .retry(&this.session_id, cx)
1948                    .map(|retry| retry.run(cx))
1949            })?
1950            .context("retrying a session is not supported")?
1951            .await
1952        })
1953    }
1954
1955    fn run_turn(
1956        &mut self,
1957        cx: &mut Context<Self>,
1958        f: impl 'static + AsyncFnOnce(WeakEntity<Self>, &mut AsyncApp) -> Result<acp::PromptResponse>,
1959    ) -> BoxFuture<'static, Result<Option<acp::PromptResponse>>> {
1960        self.clear_completed_plan_entries(cx);
1961        self.had_error = false;
1962
1963        let (tx, rx) = oneshot::channel();
1964        let cancel_task = self.cancel(cx);
1965
1966        self.turn_id += 1;
1967        let turn_id = self.turn_id;
1968        self.running_turn = Some(RunningTurn {
1969            id: turn_id,
1970            send_task: cx.spawn(async move |this, cx| {
1971                cancel_task.await;
1972                tx.send(f(this, cx).await).ok();
1973            }),
1974        });
1975
1976        cx.spawn(async move |this, cx| {
1977            let response = rx.await;
1978
1979            this.update(cx, |this, cx| this.update_last_checkpoint(cx))?
1980                .await?;
1981
1982            this.update(cx, |this, cx| {
1983                if this.parent_session_id.is_none() {
1984                    this.project
1985                        .update(cx, |project, cx| project.set_agent_location(None, cx));
1986                }
1987                let Ok(response) = response else {
1988                    // tx dropped, just return
1989                    return Ok(None);
1990                };
1991
1992                let is_same_turn = this
1993                    .running_turn
1994                    .as_ref()
1995                    .is_some_and(|turn| turn_id == turn.id);
1996
1997                // If the user submitted a follow up message, running_turn might
1998                // already point to a different turn. Therefore we only want to
1999                // take the task if it's the same turn.
2000                if is_same_turn {
2001                    this.running_turn.take();
2002                }
2003
2004                match response {
2005                    Ok(r) => {
2006                        if r.stop_reason == acp::StopReason::MaxTokens {
2007                            this.had_error = true;
2008                            cx.emit(AcpThreadEvent::Error);
2009                            log::error!("Max tokens reached. Usage: {:?}", this.token_usage);
2010                            return Err(anyhow!("Max tokens reached"));
2011                        }
2012
2013                        let canceled = matches!(r.stop_reason, acp::StopReason::Cancelled);
2014                        if canceled {
2015                            this.mark_pending_tools_as_canceled();
2016                        }
2017
2018                        // Handle refusal - distinguish between user prompt and tool call refusals
2019                        if let acp::StopReason::Refusal = r.stop_reason {
2020                            this.had_error = true;
2021                            if let Some((user_msg_ix, _)) = this.last_user_message() {
2022                                // Check if there's a completed tool call with results after the last user message
2023                                // This indicates the refusal is in response to tool output, not the user's prompt
2024                                let has_completed_tool_call_after_user_msg =
2025                                    this.entries.iter().skip(user_msg_ix + 1).any(|entry| {
2026                                        if let AgentThreadEntry::ToolCall(tool_call) = entry {
2027                                            // Check if the tool call has completed and has output
2028                                            matches!(tool_call.status, ToolCallStatus::Completed)
2029                                                && tool_call.raw_output.is_some()
2030                                        } else {
2031                                            false
2032                                        }
2033                                    });
2034
2035                                if has_completed_tool_call_after_user_msg {
2036                                    // Refusal is due to tool output - don't truncate, just notify
2037                                    // The model refused based on what the tool returned
2038                                    cx.emit(AcpThreadEvent::Refusal);
2039                                } else {
2040                                    // User prompt was refused - truncate back to before the user message
2041                                    let range = user_msg_ix..this.entries.len();
2042                                    if range.start < range.end {
2043                                        this.entries.truncate(user_msg_ix);
2044                                        cx.emit(AcpThreadEvent::EntriesRemoved(range));
2045                                    }
2046                                    cx.emit(AcpThreadEvent::Refusal);
2047                                }
2048                            } else {
2049                                // No user message found, treat as general refusal
2050                                cx.emit(AcpThreadEvent::Refusal);
2051                            }
2052                        }
2053
2054                        cx.emit(AcpThreadEvent::Stopped(r.stop_reason));
2055                        Ok(Some(r))
2056                    }
2057                    Err(e) => {
2058                        this.had_error = true;
2059                        cx.emit(AcpThreadEvent::Error);
2060                        log::error!("Error in run turn: {:?}", e);
2061                        Err(e)
2062                    }
2063                }
2064            })?
2065        })
2066        .boxed()
2067    }
2068
2069    pub fn cancel(&mut self, cx: &mut Context<Self>) -> Task<()> {
2070        let Some(turn) = self.running_turn.take() else {
2071            return Task::ready(());
2072        };
2073        self.connection.cancel(&self.session_id, cx);
2074
2075        self.mark_pending_tools_as_canceled();
2076
2077        // Wait for the send task to complete
2078        cx.background_spawn(turn.send_task)
2079    }
2080
2081    fn mark_pending_tools_as_canceled(&mut self) {
2082        for entry in self.entries.iter_mut() {
2083            if let AgentThreadEntry::ToolCall(call) = entry {
2084                let cancel = matches!(
2085                    call.status,
2086                    ToolCallStatus::Pending
2087                        | ToolCallStatus::WaitingForConfirmation { .. }
2088                        | ToolCallStatus::InProgress
2089                );
2090
2091                if cancel {
2092                    call.status = ToolCallStatus::Canceled;
2093                }
2094            }
2095        }
2096    }
2097
2098    /// Restores the git working tree to the state at the given checkpoint (if one exists)
2099    pub fn restore_checkpoint(
2100        &mut self,
2101        id: UserMessageId,
2102        cx: &mut Context<Self>,
2103    ) -> Task<Result<()>> {
2104        let Some((_, message)) = self.user_message_mut(&id) else {
2105            return Task::ready(Err(anyhow!("message not found")));
2106        };
2107
2108        let checkpoint = message
2109            .checkpoint
2110            .as_ref()
2111            .map(|c| c.git_checkpoint.clone());
2112
2113        // Cancel any in-progress generation before restoring
2114        let cancel_task = self.cancel(cx);
2115        let rewind = self.rewind(id.clone(), cx);
2116        let git_store = self.project.read(cx).git_store().clone();
2117
2118        cx.spawn(async move |_, cx| {
2119            cancel_task.await;
2120            rewind.await?;
2121            if let Some(checkpoint) = checkpoint {
2122                git_store
2123                    .update(cx, |git, cx| git.restore_checkpoint(checkpoint, cx))
2124                    .await?;
2125            }
2126
2127            Ok(())
2128        })
2129    }
2130
2131    /// Rewinds this thread to before the entry at `index`, removing it and all
2132    /// subsequent entries while rejecting any action_log changes made from that point.
2133    /// Unlike `restore_checkpoint`, this method does not restore from git.
2134    pub fn rewind(&mut self, id: UserMessageId, cx: &mut Context<Self>) -> Task<Result<()>> {
2135        let Some(truncate) = self.connection.truncate(&self.session_id, cx) else {
2136            return Task::ready(Err(anyhow!("not supported")));
2137        };
2138
2139        let telemetry = ActionLogTelemetry::from(&*self);
2140        cx.spawn(async move |this, cx| {
2141            cx.update(|cx| truncate.run(id.clone(), cx)).await?;
2142            this.update(cx, |this, cx| {
2143                if let Some((ix, _)) = this.user_message_mut(&id) {
2144                    // Collect all terminals from entries that will be removed
2145                    let terminals_to_remove: Vec<acp::TerminalId> = this.entries[ix..]
2146                        .iter()
2147                        .flat_map(|entry| entry.terminals())
2148                        .filter_map(|terminal| terminal.read(cx).id().clone().into())
2149                        .collect();
2150
2151                    let range = ix..this.entries.len();
2152                    this.entries.truncate(ix);
2153                    cx.emit(AcpThreadEvent::EntriesRemoved(range));
2154
2155                    // Kill and remove the terminals
2156                    for terminal_id in terminals_to_remove {
2157                        if let Some(terminal) = this.terminals.remove(&terminal_id) {
2158                            terminal.update(cx, |terminal, cx| {
2159                                terminal.kill(cx);
2160                            });
2161                        }
2162                    }
2163                }
2164                this.action_log().update(cx, |action_log, cx| {
2165                    action_log.reject_all_edits(Some(telemetry), cx)
2166                })
2167            })?
2168            .await;
2169            Ok(())
2170        })
2171    }
2172
2173    fn update_last_checkpoint(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
2174        let git_store = self.project.read(cx).git_store().clone();
2175
2176        let Some((_, message)) = self.last_user_message() else {
2177            return Task::ready(Ok(()));
2178        };
2179        let Some(user_message_id) = message.id.clone() else {
2180            return Task::ready(Ok(()));
2181        };
2182        let Some(checkpoint) = message.checkpoint.as_ref() else {
2183            return Task::ready(Ok(()));
2184        };
2185        let old_checkpoint = checkpoint.git_checkpoint.clone();
2186
2187        let new_checkpoint = git_store.update(cx, |git, cx| git.checkpoint(cx));
2188        cx.spawn(async move |this, cx| {
2189            let Some(new_checkpoint) = new_checkpoint
2190                .await
2191                .context("failed to get new checkpoint")
2192                .log_err()
2193            else {
2194                return Ok(());
2195            };
2196
2197            let equal = git_store
2198                .update(cx, |git, cx| {
2199                    git.compare_checkpoints(old_checkpoint.clone(), new_checkpoint, cx)
2200                })
2201                .await
2202                .unwrap_or(true);
2203
2204            this.update(cx, |this, cx| {
2205                if let Some((ix, message)) = this.user_message_mut(&user_message_id) {
2206                    if let Some(checkpoint) = message.checkpoint.as_mut() {
2207                        checkpoint.show = !equal;
2208                        cx.emit(AcpThreadEvent::EntryUpdated(ix));
2209                    }
2210                }
2211            })?;
2212
2213            Ok(())
2214        })
2215    }
2216
2217    fn last_user_message(&mut self) -> Option<(usize, &mut UserMessage)> {
2218        self.entries
2219            .iter_mut()
2220            .enumerate()
2221            .rev()
2222            .find_map(|(ix, entry)| {
2223                if let AgentThreadEntry::UserMessage(message) = entry {
2224                    Some((ix, message))
2225                } else {
2226                    None
2227                }
2228            })
2229    }
2230
2231    fn user_message_mut(&mut self, id: &UserMessageId) -> Option<(usize, &mut UserMessage)> {
2232        self.entries.iter_mut().enumerate().find_map(|(ix, entry)| {
2233            if let AgentThreadEntry::UserMessage(message) = entry {
2234                if message.id.as_ref() == Some(id) {
2235                    Some((ix, message))
2236                } else {
2237                    None
2238                }
2239            } else {
2240                None
2241            }
2242        })
2243    }
2244
2245    pub fn read_text_file(
2246        &self,
2247        path: PathBuf,
2248        line: Option<u32>,
2249        limit: Option<u32>,
2250        reuse_shared_snapshot: bool,
2251        cx: &mut Context<Self>,
2252    ) -> Task<Result<String, acp::Error>> {
2253        // Args are 1-based, move to 0-based
2254        let line = line.unwrap_or_default().saturating_sub(1);
2255        let limit = limit.unwrap_or(u32::MAX);
2256        let project = self.project.clone();
2257        let action_log = self.action_log.clone();
2258        let should_update_agent_location = self.parent_session_id.is_none();
2259        cx.spawn(async move |this, cx| {
2260            let load = project.update(cx, |project, cx| {
2261                let path = project
2262                    .project_path_for_absolute_path(&path, cx)
2263                    .ok_or_else(|| {
2264                        acp::Error::resource_not_found(Some(path.display().to_string()))
2265                    })?;
2266                Ok::<_, acp::Error>(project.open_buffer(path, cx))
2267            })?;
2268
2269            let buffer = load.await?;
2270
2271            let snapshot = if reuse_shared_snapshot {
2272                this.read_with(cx, |this, _| {
2273                    this.shared_buffers.get(&buffer.clone()).cloned()
2274                })
2275                .log_err()
2276                .flatten()
2277            } else {
2278                None
2279            };
2280
2281            let snapshot = if let Some(snapshot) = snapshot {
2282                snapshot
2283            } else {
2284                action_log.update(cx, |action_log, cx| {
2285                    action_log.buffer_read(buffer.clone(), cx);
2286                });
2287
2288                let snapshot = buffer.update(cx, |buffer, _| buffer.snapshot());
2289                this.update(cx, |this, _| {
2290                    this.shared_buffers.insert(buffer.clone(), snapshot.clone());
2291                })?;
2292                snapshot
2293            };
2294
2295            let max_point = snapshot.max_point();
2296            let start_position = Point::new(line, 0);
2297
2298            if start_position > max_point {
2299                return Err(acp::Error::invalid_params().data(format!(
2300                    "Attempting to read beyond the end of the file, line {}:{}",
2301                    max_point.row + 1,
2302                    max_point.column
2303                )));
2304            }
2305
2306            let start = snapshot.anchor_before(start_position);
2307            let end = snapshot.anchor_before(Point::new(line.saturating_add(limit), 0));
2308
2309            if should_update_agent_location {
2310                project.update(cx, |project, cx| {
2311                    project.set_agent_location(
2312                        Some(AgentLocation {
2313                            buffer: buffer.downgrade(),
2314                            position: start,
2315                        }),
2316                        cx,
2317                    );
2318                });
2319            }
2320
2321            Ok(snapshot.text_for_range(start..end).collect::<String>())
2322        })
2323    }
2324
2325    pub fn write_text_file(
2326        &self,
2327        path: PathBuf,
2328        content: String,
2329        cx: &mut Context<Self>,
2330    ) -> Task<Result<()>> {
2331        let project = self.project.clone();
2332        let action_log = self.action_log.clone();
2333        let should_update_agent_location = self.parent_session_id.is_none();
2334        cx.spawn(async move |this, cx| {
2335            let load = project.update(cx, |project, cx| {
2336                let path = project
2337                    .project_path_for_absolute_path(&path, cx)
2338                    .context("invalid path")?;
2339                anyhow::Ok(project.open_buffer(path, cx))
2340            });
2341            let buffer = load?.await?;
2342            let snapshot = this.update(cx, |this, cx| {
2343                this.shared_buffers
2344                    .get(&buffer)
2345                    .cloned()
2346                    .unwrap_or_else(|| buffer.read(cx).snapshot())
2347            })?;
2348            let edits = cx
2349                .background_executor()
2350                .spawn(async move {
2351                    let old_text = snapshot.text();
2352                    text_diff(old_text.as_str(), &content)
2353                        .into_iter()
2354                        .map(|(range, replacement)| {
2355                            (snapshot.anchor_range_around(range), replacement)
2356                        })
2357                        .collect::<Vec<_>>()
2358                })
2359                .await;
2360
2361            if should_update_agent_location {
2362                project.update(cx, |project, cx| {
2363                    project.set_agent_location(
2364                        Some(AgentLocation {
2365                            buffer: buffer.downgrade(),
2366                            position: edits
2367                                .last()
2368                                .map(|(range, _)| range.end)
2369                                .unwrap_or(Anchor::min_for_buffer(buffer.read(cx).remote_id())),
2370                        }),
2371                        cx,
2372                    );
2373                });
2374            }
2375
2376            let format_on_save = cx.update(|cx| {
2377                action_log.update(cx, |action_log, cx| {
2378                    action_log.buffer_read(buffer.clone(), cx);
2379                });
2380
2381                let format_on_save = buffer.update(cx, |buffer, cx| {
2382                    buffer.edit(edits, None, cx);
2383
2384                    let settings = language::language_settings::language_settings(
2385                        buffer.language().map(|l| l.name()),
2386                        buffer.file(),
2387                        cx,
2388                    );
2389
2390                    settings.format_on_save != FormatOnSave::Off
2391                });
2392                action_log.update(cx, |action_log, cx| {
2393                    action_log.buffer_edited(buffer.clone(), cx);
2394                });
2395                format_on_save
2396            });
2397
2398            if format_on_save {
2399                let format_task = project.update(cx, |project, cx| {
2400                    project.format(
2401                        HashSet::from_iter([buffer.clone()]),
2402                        LspFormatTarget::Buffers,
2403                        false,
2404                        FormatTrigger::Save,
2405                        cx,
2406                    )
2407                });
2408                format_task.await.log_err();
2409
2410                action_log.update(cx, |action_log, cx| {
2411                    action_log.buffer_edited(buffer.clone(), cx);
2412                });
2413            }
2414
2415            project
2416                .update(cx, |project, cx| project.save_buffer(buffer, cx))
2417                .await
2418        })
2419    }
2420
2421    pub fn create_terminal(
2422        &self,
2423        command: String,
2424        args: Vec<String>,
2425        extra_env: Vec<acp::EnvVariable>,
2426        cwd: Option<PathBuf>,
2427        output_byte_limit: Option<u64>,
2428        cx: &mut Context<Self>,
2429    ) -> Task<Result<Entity<Terminal>>> {
2430        let env = match &cwd {
2431            Some(dir) => self.project.update(cx, |project, cx| {
2432                project.environment().update(cx, |env, cx| {
2433                    env.directory_environment(dir.as_path().into(), cx)
2434                })
2435            }),
2436            None => Task::ready(None).shared(),
2437        };
2438        let env = cx.spawn(async move |_, _| {
2439            let mut env = env.await.unwrap_or_default();
2440            // Disables paging for `git` and hopefully other commands
2441            env.insert("PAGER".into(), "".into());
2442            for var in extra_env {
2443                env.insert(var.name, var.value);
2444            }
2445            env
2446        });
2447
2448        let project = self.project.clone();
2449        let language_registry = project.read(cx).languages().clone();
2450        let is_windows = project.read(cx).path_style(cx).is_windows();
2451
2452        let terminal_id = acp::TerminalId::new(Uuid::new_v4().to_string());
2453        let terminal_task = cx.spawn({
2454            let terminal_id = terminal_id.clone();
2455            async move |_this, cx| {
2456                let env = env.await;
2457                let shell = project
2458                    .update(cx, |project, cx| {
2459                        project
2460                            .remote_client()
2461                            .and_then(|r| r.read(cx).default_system_shell())
2462                    })
2463                    .unwrap_or_else(|| get_default_system_shell_preferring_bash());
2464                let (task_command, task_args) =
2465                    ShellBuilder::new(&Shell::Program(shell), is_windows)
2466                        .redirect_stdin_to_dev_null()
2467                        .build(Some(command.clone()), &args);
2468                let terminal = project
2469                    .update(cx, |project, cx| {
2470                        project.create_terminal_task(
2471                            task::SpawnInTerminal {
2472                                command: Some(task_command),
2473                                args: task_args,
2474                                cwd: cwd.clone(),
2475                                env,
2476                                ..Default::default()
2477                            },
2478                            cx,
2479                        )
2480                    })
2481                    .await?;
2482
2483                anyhow::Ok(cx.new(|cx| {
2484                    Terminal::new(
2485                        terminal_id,
2486                        &format!("{} {}", command, args.join(" ")),
2487                        cwd,
2488                        output_byte_limit.map(|l| l as usize),
2489                        terminal,
2490                        language_registry,
2491                        cx,
2492                    )
2493                }))
2494            }
2495        });
2496
2497        cx.spawn(async move |this, cx| {
2498            let terminal = terminal_task.await?;
2499            this.update(cx, |this, _cx| {
2500                this.terminals.insert(terminal_id, terminal.clone());
2501                terminal
2502            })
2503        })
2504    }
2505
2506    pub fn kill_terminal(
2507        &mut self,
2508        terminal_id: acp::TerminalId,
2509        cx: &mut Context<Self>,
2510    ) -> Result<()> {
2511        self.terminals
2512            .get(&terminal_id)
2513            .context("Terminal not found")?
2514            .update(cx, |terminal, cx| {
2515                terminal.kill(cx);
2516            });
2517
2518        Ok(())
2519    }
2520
2521    pub fn release_terminal(
2522        &mut self,
2523        terminal_id: acp::TerminalId,
2524        cx: &mut Context<Self>,
2525    ) -> Result<()> {
2526        self.terminals
2527            .remove(&terminal_id)
2528            .context("Terminal not found")?
2529            .update(cx, |terminal, cx| {
2530                terminal.kill(cx);
2531            });
2532
2533        Ok(())
2534    }
2535
2536    pub fn terminal(&self, terminal_id: acp::TerminalId) -> Result<Entity<Terminal>> {
2537        self.terminals
2538            .get(&terminal_id)
2539            .context("Terminal not found")
2540            .cloned()
2541    }
2542
2543    pub fn to_markdown(&self, cx: &App) -> String {
2544        self.entries.iter().map(|e| e.to_markdown(cx)).collect()
2545    }
2546
2547    pub fn emit_load_error(&mut self, error: LoadError, cx: &mut Context<Self>) {
2548        cx.emit(AcpThreadEvent::LoadError(error));
2549    }
2550
2551    pub fn register_terminal_created(
2552        &mut self,
2553        terminal_id: acp::TerminalId,
2554        command_label: String,
2555        working_dir: Option<PathBuf>,
2556        output_byte_limit: Option<u64>,
2557        terminal: Entity<::terminal::Terminal>,
2558        cx: &mut Context<Self>,
2559    ) -> Entity<Terminal> {
2560        let language_registry = self.project.read(cx).languages().clone();
2561
2562        let entity = cx.new(|cx| {
2563            Terminal::new(
2564                terminal_id.clone(),
2565                &command_label,
2566                working_dir.clone(),
2567                output_byte_limit.map(|l| l as usize),
2568                terminal,
2569                language_registry,
2570                cx,
2571            )
2572        });
2573        self.terminals.insert(terminal_id.clone(), entity.clone());
2574        entity
2575    }
2576
2577    pub fn mark_as_subagent_output(&mut self, cx: &mut Context<Self>) {
2578        for entry in self.entries.iter_mut().rev() {
2579            if let AgentThreadEntry::AssistantMessage(assistant_message) = entry {
2580                assistant_message.is_subagent_output = true;
2581                cx.notify();
2582                return;
2583            }
2584        }
2585    }
2586}
2587
2588fn markdown_for_raw_output(
2589    raw_output: &serde_json::Value,
2590    language_registry: &Arc<LanguageRegistry>,
2591    cx: &mut App,
2592) -> Option<Entity<Markdown>> {
2593    match raw_output {
2594        serde_json::Value::Null => None,
2595        serde_json::Value::Bool(value) => Some(cx.new(|cx| {
2596            Markdown::new(
2597                value.to_string().into(),
2598                Some(language_registry.clone()),
2599                None,
2600                cx,
2601            )
2602        })),
2603        serde_json::Value::Number(value) => Some(cx.new(|cx| {
2604            Markdown::new(
2605                value.to_string().into(),
2606                Some(language_registry.clone()),
2607                None,
2608                cx,
2609            )
2610        })),
2611        serde_json::Value::String(value) => Some(cx.new(|cx| {
2612            Markdown::new(
2613                value.clone().into(),
2614                Some(language_registry.clone()),
2615                None,
2616                cx,
2617            )
2618        })),
2619        value => Some(cx.new(|cx| {
2620            let pretty_json = to_string_pretty(value).unwrap_or_else(|_| value.to_string());
2621
2622            Markdown::new(
2623                format!("```json\n{}\n```", pretty_json).into(),
2624                Some(language_registry.clone()),
2625                None,
2626                cx,
2627            )
2628        })),
2629    }
2630}
2631
2632#[cfg(test)]
2633mod tests {
2634    use super::*;
2635    use anyhow::anyhow;
2636    use futures::{channel::mpsc, future::LocalBoxFuture, select};
2637    use gpui::{App, AsyncApp, TestAppContext, WeakEntity};
2638    use indoc::indoc;
2639    use project::{FakeFs, Fs};
2640    use rand::{distr, prelude::*};
2641    use serde_json::json;
2642    use settings::SettingsStore;
2643    use smol::stream::StreamExt as _;
2644    use std::{
2645        any::Any,
2646        cell::RefCell,
2647        path::Path,
2648        rc::Rc,
2649        sync::atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
2650        time::Duration,
2651    };
2652    use util::path;
2653
2654    fn init_test(cx: &mut TestAppContext) {
2655        env_logger::try_init().ok();
2656        cx.update(|cx| {
2657            let settings_store = SettingsStore::test(cx);
2658            cx.set_global(settings_store);
2659        });
2660    }
2661
2662    #[gpui::test]
2663    async fn test_terminal_output_buffered_before_created_renders(cx: &mut gpui::TestAppContext) {
2664        init_test(cx);
2665
2666        let fs = FakeFs::new(cx.executor());
2667        let project = Project::test(fs, [], cx).await;
2668        let connection = Rc::new(FakeAgentConnection::new());
2669        let thread = cx
2670            .update(|cx| connection.new_session(project, std::path::Path::new(path!("/test")), cx))
2671            .await
2672            .unwrap();
2673
2674        let terminal_id = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
2675
2676        // Send Output BEFORE Created - should be buffered by acp_thread
2677        thread.update(cx, |thread, cx| {
2678            thread.on_terminal_provider_event(
2679                TerminalProviderEvent::Output {
2680                    terminal_id: terminal_id.clone(),
2681                    data: b"hello buffered".to_vec(),
2682                },
2683                cx,
2684            );
2685        });
2686
2687        // Create a display-only terminal and then send Created
2688        let lower = cx.new(|cx| {
2689            let builder = ::terminal::TerminalBuilder::new_display_only(
2690                ::terminal::terminal_settings::CursorShape::default(),
2691                ::terminal::terminal_settings::AlternateScroll::On,
2692                None,
2693                0,
2694                cx.background_executor(),
2695                PathStyle::local(),
2696            )
2697            .unwrap();
2698            builder.subscribe(cx)
2699        });
2700
2701        thread.update(cx, |thread, cx| {
2702            thread.on_terminal_provider_event(
2703                TerminalProviderEvent::Created {
2704                    terminal_id: terminal_id.clone(),
2705                    label: "Buffered Test".to_string(),
2706                    cwd: None,
2707                    output_byte_limit: None,
2708                    terminal: lower.clone(),
2709                },
2710                cx,
2711            );
2712        });
2713
2714        // After Created, buffered Output should have been flushed into the renderer
2715        let content = thread.read_with(cx, |thread, cx| {
2716            let term = thread.terminal(terminal_id.clone()).unwrap();
2717            term.read_with(cx, |t, cx| t.inner().read(cx).get_content())
2718        });
2719
2720        assert!(
2721            content.contains("hello buffered"),
2722            "expected buffered output to render, got: {content}"
2723        );
2724    }
2725
2726    #[gpui::test]
2727    async fn test_terminal_output_and_exit_buffered_before_created(cx: &mut gpui::TestAppContext) {
2728        init_test(cx);
2729
2730        let fs = FakeFs::new(cx.executor());
2731        let project = Project::test(fs, [], cx).await;
2732        let connection = Rc::new(FakeAgentConnection::new());
2733        let thread = cx
2734            .update(|cx| connection.new_session(project, std::path::Path::new(path!("/test")), cx))
2735            .await
2736            .unwrap();
2737
2738        let terminal_id = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
2739
2740        // Send Output BEFORE Created
2741        thread.update(cx, |thread, cx| {
2742            thread.on_terminal_provider_event(
2743                TerminalProviderEvent::Output {
2744                    terminal_id: terminal_id.clone(),
2745                    data: b"pre-exit data".to_vec(),
2746                },
2747                cx,
2748            );
2749        });
2750
2751        // Send Exit BEFORE Created
2752        thread.update(cx, |thread, cx| {
2753            thread.on_terminal_provider_event(
2754                TerminalProviderEvent::Exit {
2755                    terminal_id: terminal_id.clone(),
2756                    status: acp::TerminalExitStatus::new().exit_code(0),
2757                },
2758                cx,
2759            );
2760        });
2761
2762        // Now create a display-only lower-level terminal and send Created
2763        let lower = cx.new(|cx| {
2764            let builder = ::terminal::TerminalBuilder::new_display_only(
2765                ::terminal::terminal_settings::CursorShape::default(),
2766                ::terminal::terminal_settings::AlternateScroll::On,
2767                None,
2768                0,
2769                cx.background_executor(),
2770                PathStyle::local(),
2771            )
2772            .unwrap();
2773            builder.subscribe(cx)
2774        });
2775
2776        thread.update(cx, |thread, cx| {
2777            thread.on_terminal_provider_event(
2778                TerminalProviderEvent::Created {
2779                    terminal_id: terminal_id.clone(),
2780                    label: "Buffered Exit Test".to_string(),
2781                    cwd: None,
2782                    output_byte_limit: None,
2783                    terminal: lower.clone(),
2784                },
2785                cx,
2786            );
2787        });
2788
2789        // Output should be present after Created (flushed from buffer)
2790        let content = thread.read_with(cx, |thread, cx| {
2791            let term = thread.terminal(terminal_id.clone()).unwrap();
2792            term.read_with(cx, |t, cx| t.inner().read(cx).get_content())
2793        });
2794
2795        assert!(
2796            content.contains("pre-exit data"),
2797            "expected pre-exit data to render, got: {content}"
2798        );
2799    }
2800
2801    /// Test that killing a terminal via Terminal::kill properly:
2802    /// 1. Causes wait_for_exit to complete (doesn't hang forever)
2803    /// 2. The underlying terminal still has the output that was written before the kill
2804    ///
2805    /// This test verifies that the fix to kill_active_task (which now also kills
2806    /// the shell process in addition to the foreground process) properly allows
2807    /// wait_for_exit to complete instead of hanging indefinitely.
2808    #[cfg(unix)]
2809    #[gpui::test]
2810    async fn test_terminal_kill_allows_wait_for_exit_to_complete(cx: &mut gpui::TestAppContext) {
2811        use std::collections::HashMap;
2812        use task::Shell;
2813        use util::shell_builder::ShellBuilder;
2814
2815        init_test(cx);
2816        cx.executor().allow_parking();
2817
2818        let fs = FakeFs::new(cx.executor());
2819        let project = Project::test(fs, [], cx).await;
2820        let connection = Rc::new(FakeAgentConnection::new());
2821        let thread = cx
2822            .update(|cx| connection.new_session(project.clone(), Path::new(path!("/test")), cx))
2823            .await
2824            .unwrap();
2825
2826        let terminal_id = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
2827
2828        // Create a real PTY terminal that runs a command which prints output then sleeps
2829        // We use printf instead of echo and chain with && sleep to ensure proper execution
2830        let (completion_tx, _completion_rx) = smol::channel::unbounded();
2831        let (program, args) = ShellBuilder::new(&Shell::System, false).build(
2832            Some("printf 'output_before_kill\\n' && sleep 60".to_owned()),
2833            &[],
2834        );
2835
2836        let builder = cx
2837            .update(|cx| {
2838                ::terminal::TerminalBuilder::new(
2839                    None,
2840                    None,
2841                    task::Shell::WithArguments {
2842                        program,
2843                        args,
2844                        title_override: None,
2845                    },
2846                    HashMap::default(),
2847                    ::terminal::terminal_settings::CursorShape::default(),
2848                    ::terminal::terminal_settings::AlternateScroll::On,
2849                    None,
2850                    vec![],
2851                    0,
2852                    false,
2853                    0,
2854                    Some(completion_tx),
2855                    cx,
2856                    vec![],
2857                    PathStyle::local(),
2858                )
2859            })
2860            .await
2861            .unwrap();
2862
2863        let lower_terminal = cx.new(|cx| builder.subscribe(cx));
2864
2865        // Create the acp_thread Terminal wrapper
2866        thread.update(cx, |thread, cx| {
2867            thread.on_terminal_provider_event(
2868                TerminalProviderEvent::Created {
2869                    terminal_id: terminal_id.clone(),
2870                    label: "printf output_before_kill && sleep 60".to_string(),
2871                    cwd: None,
2872                    output_byte_limit: None,
2873                    terminal: lower_terminal.clone(),
2874                },
2875                cx,
2876            );
2877        });
2878
2879        // Wait for the printf command to execute and produce output
2880        // Use real time since parking is enabled
2881        cx.executor().timer(Duration::from_millis(500)).await;
2882
2883        // Get the acp_thread Terminal and kill it
2884        let wait_for_exit = thread.update(cx, |thread, cx| {
2885            let term = thread.terminals.get(&terminal_id).unwrap();
2886            let wait_for_exit = term.read(cx).wait_for_exit();
2887            term.update(cx, |term, cx| {
2888                term.kill(cx);
2889            });
2890            wait_for_exit
2891        });
2892
2893        // KEY ASSERTION: wait_for_exit should complete within a reasonable time (not hang).
2894        // Before the fix to kill_active_task, this would hang forever because
2895        // only the foreground process was killed, not the shell, so the PTY
2896        // child never exited and wait_for_completed_task never completed.
2897        let exit_result = futures::select! {
2898            result = futures::FutureExt::fuse(wait_for_exit) => Some(result),
2899            _ = futures::FutureExt::fuse(cx.background_executor.timer(Duration::from_secs(5))) => None,
2900        };
2901
2902        assert!(
2903            exit_result.is_some(),
2904            "wait_for_exit should complete after kill, but it timed out. \
2905            This indicates kill_active_task is not properly killing the shell process."
2906        );
2907
2908        // Give the system a chance to process any pending updates
2909        cx.run_until_parked();
2910
2911        // Verify that the underlying terminal still has the output that was
2912        // written before the kill. This verifies that killing doesn't lose output.
2913        let inner_content = thread.read_with(cx, |thread, cx| {
2914            let term = thread.terminals.get(&terminal_id).unwrap();
2915            term.read(cx).inner().read(cx).get_content()
2916        });
2917
2918        assert!(
2919            inner_content.contains("output_before_kill"),
2920            "Underlying terminal should contain output from before kill, got: {}",
2921            inner_content
2922        );
2923    }
2924
2925    #[gpui::test]
2926    async fn test_push_user_content_block(cx: &mut gpui::TestAppContext) {
2927        init_test(cx);
2928
2929        let fs = FakeFs::new(cx.executor());
2930        let project = Project::test(fs, [], cx).await;
2931        let connection = Rc::new(FakeAgentConnection::new());
2932        let thread = cx
2933            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
2934            .await
2935            .unwrap();
2936
2937        // Test creating a new user message
2938        thread.update(cx, |thread, cx| {
2939            thread.push_user_content_block(None, "Hello, ".into(), cx);
2940        });
2941
2942        thread.update(cx, |thread, cx| {
2943            assert_eq!(thread.entries.len(), 1);
2944            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
2945                assert_eq!(user_msg.id, None);
2946                assert_eq!(user_msg.content.to_markdown(cx), "Hello, ");
2947            } else {
2948                panic!("Expected UserMessage");
2949            }
2950        });
2951
2952        // Test appending to existing user message
2953        let message_1_id = UserMessageId::new();
2954        thread.update(cx, |thread, cx| {
2955            thread.push_user_content_block(Some(message_1_id.clone()), "world!".into(), cx);
2956        });
2957
2958        thread.update(cx, |thread, cx| {
2959            assert_eq!(thread.entries.len(), 1);
2960            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[0] {
2961                assert_eq!(user_msg.id, Some(message_1_id));
2962                assert_eq!(user_msg.content.to_markdown(cx), "Hello, world!");
2963            } else {
2964                panic!("Expected UserMessage");
2965            }
2966        });
2967
2968        // Test creating new user message after assistant message
2969        thread.update(cx, |thread, cx| {
2970            thread.push_assistant_content_block("Assistant response".into(), false, cx);
2971        });
2972
2973        let message_2_id = UserMessageId::new();
2974        thread.update(cx, |thread, cx| {
2975            thread.push_user_content_block(
2976                Some(message_2_id.clone()),
2977                "New user message".into(),
2978                cx,
2979            );
2980        });
2981
2982        thread.update(cx, |thread, cx| {
2983            assert_eq!(thread.entries.len(), 3);
2984            if let AgentThreadEntry::UserMessage(user_msg) = &thread.entries[2] {
2985                assert_eq!(user_msg.id, Some(message_2_id));
2986                assert_eq!(user_msg.content.to_markdown(cx), "New user message");
2987            } else {
2988                panic!("Expected UserMessage at index 2");
2989            }
2990        });
2991    }
2992
2993    #[gpui::test]
2994    async fn test_thinking_concatenation(cx: &mut gpui::TestAppContext) {
2995        init_test(cx);
2996
2997        let fs = FakeFs::new(cx.executor());
2998        let project = Project::test(fs, [], cx).await;
2999        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
3000            |_, thread, mut cx| {
3001                async move {
3002                    thread.update(&mut cx, |thread, cx| {
3003                        thread
3004                            .handle_session_update(
3005                                acp::SessionUpdate::AgentThoughtChunk(acp::ContentChunk::new(
3006                                    "Thinking ".into(),
3007                                )),
3008                                cx,
3009                            )
3010                            .unwrap();
3011                        thread
3012                            .handle_session_update(
3013                                acp::SessionUpdate::AgentThoughtChunk(acp::ContentChunk::new(
3014                                    "hard!".into(),
3015                                )),
3016                                cx,
3017                            )
3018                            .unwrap();
3019                    })?;
3020                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3021                }
3022                .boxed_local()
3023            },
3024        ));
3025
3026        let thread = cx
3027            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3028            .await
3029            .unwrap();
3030
3031        thread
3032            .update(cx, |thread, cx| thread.send_raw("Hello from Zed!", cx))
3033            .await
3034            .unwrap();
3035
3036        let output = thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
3037        assert_eq!(
3038            output,
3039            indoc! {r#"
3040            ## User
3041
3042            Hello from Zed!
3043
3044            ## Assistant
3045
3046            <thinking>
3047            Thinking hard!
3048            </thinking>
3049
3050            "#}
3051        );
3052    }
3053
3054    #[gpui::test]
3055    async fn test_edits_concurrently_to_user(cx: &mut TestAppContext) {
3056        init_test(cx);
3057
3058        let fs = FakeFs::new(cx.executor());
3059        fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\n"}))
3060            .await;
3061        let project = Project::test(fs.clone(), [], cx).await;
3062        let (read_file_tx, read_file_rx) = oneshot::channel::<()>();
3063        let read_file_tx = Rc::new(RefCell::new(Some(read_file_tx)));
3064        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
3065            move |_, thread, mut cx| {
3066                let read_file_tx = read_file_tx.clone();
3067                async move {
3068                    let content = thread
3069                        .update(&mut cx, |thread, cx| {
3070                            thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
3071                        })
3072                        .unwrap()
3073                        .await
3074                        .unwrap();
3075                    assert_eq!(content, "one\ntwo\nthree\n");
3076                    read_file_tx.take().unwrap().send(()).unwrap();
3077                    thread
3078                        .update(&mut cx, |thread, cx| {
3079                            thread.write_text_file(
3080                                path!("/tmp/foo").into(),
3081                                "one\ntwo\nthree\nfour\nfive\n".to_string(),
3082                                cx,
3083                            )
3084                        })
3085                        .unwrap()
3086                        .await
3087                        .unwrap();
3088                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3089                }
3090                .boxed_local()
3091            },
3092        ));
3093
3094        let (worktree, pathbuf) = project
3095            .update(cx, |project, cx| {
3096                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
3097            })
3098            .await
3099            .unwrap();
3100        let buffer = project
3101            .update(cx, |project, cx| {
3102                project.open_buffer((worktree.read(cx).id(), pathbuf), cx)
3103            })
3104            .await
3105            .unwrap();
3106
3107        let thread = cx
3108            .update(|cx| connection.new_session(project, Path::new(path!("/tmp")), cx))
3109            .await
3110            .unwrap();
3111
3112        let request = thread.update(cx, |thread, cx| {
3113            thread.send_raw("Extend the count in /tmp/foo", cx)
3114        });
3115        read_file_rx.await.ok();
3116        buffer.update(cx, |buffer, cx| {
3117            buffer.edit([(0..0, "zero\n".to_string())], None, cx);
3118        });
3119        cx.run_until_parked();
3120        assert_eq!(
3121            buffer.read_with(cx, |buffer, _| buffer.text()),
3122            "zero\none\ntwo\nthree\nfour\nfive\n"
3123        );
3124        assert_eq!(
3125            String::from_utf8(fs.read_file_sync(path!("/tmp/foo")).unwrap()).unwrap(),
3126            "zero\none\ntwo\nthree\nfour\nfive\n"
3127        );
3128        request.await.unwrap();
3129    }
3130
3131    #[gpui::test]
3132    async fn test_reading_from_line(cx: &mut TestAppContext) {
3133        init_test(cx);
3134
3135        let fs = FakeFs::new(cx.executor());
3136        fs.insert_tree(path!("/tmp"), json!({"foo": "one\ntwo\nthree\nfour\n"}))
3137            .await;
3138        let project = Project::test(fs.clone(), [], cx).await;
3139        project
3140            .update(cx, |project, cx| {
3141                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
3142            })
3143            .await
3144            .unwrap();
3145
3146        let connection = Rc::new(FakeAgentConnection::new());
3147
3148        let thread = cx
3149            .update(|cx| connection.new_session(project, Path::new(path!("/tmp")), cx))
3150            .await
3151            .unwrap();
3152
3153        // Whole file
3154        let content = thread
3155            .update(cx, |thread, cx| {
3156                thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
3157            })
3158            .await
3159            .unwrap();
3160
3161        assert_eq!(content, "one\ntwo\nthree\nfour\n");
3162
3163        // Only start line
3164        let content = thread
3165            .update(cx, |thread, cx| {
3166                thread.read_text_file(path!("/tmp/foo").into(), Some(3), None, false, cx)
3167            })
3168            .await
3169            .unwrap();
3170
3171        assert_eq!(content, "three\nfour\n");
3172
3173        // Only limit
3174        let content = thread
3175            .update(cx, |thread, cx| {
3176                thread.read_text_file(path!("/tmp/foo").into(), None, Some(2), false, cx)
3177            })
3178            .await
3179            .unwrap();
3180
3181        assert_eq!(content, "one\ntwo\n");
3182
3183        // Range
3184        let content = thread
3185            .update(cx, |thread, cx| {
3186                thread.read_text_file(path!("/tmp/foo").into(), Some(2), Some(2), false, cx)
3187            })
3188            .await
3189            .unwrap();
3190
3191        assert_eq!(content, "two\nthree\n");
3192
3193        // Invalid
3194        let err = thread
3195            .update(cx, |thread, cx| {
3196                thread.read_text_file(path!("/tmp/foo").into(), Some(6), Some(2), false, cx)
3197            })
3198            .await
3199            .unwrap_err();
3200
3201        assert_eq!(
3202            err.to_string(),
3203            "Invalid params: \"Attempting to read beyond the end of the file, line 5:0\""
3204        );
3205    }
3206
3207    #[gpui::test]
3208    async fn test_reading_empty_file(cx: &mut TestAppContext) {
3209        init_test(cx);
3210
3211        let fs = FakeFs::new(cx.executor());
3212        fs.insert_tree(path!("/tmp"), json!({"foo": ""})).await;
3213        let project = Project::test(fs.clone(), [], cx).await;
3214        project
3215            .update(cx, |project, cx| {
3216                project.find_or_create_worktree(path!("/tmp/foo"), true, cx)
3217            })
3218            .await
3219            .unwrap();
3220
3221        let connection = Rc::new(FakeAgentConnection::new());
3222
3223        let thread = cx
3224            .update(|cx| connection.new_session(project, Path::new(path!("/tmp")), cx))
3225            .await
3226            .unwrap();
3227
3228        // Whole file
3229        let content = thread
3230            .update(cx, |thread, cx| {
3231                thread.read_text_file(path!("/tmp/foo").into(), None, None, false, cx)
3232            })
3233            .await
3234            .unwrap();
3235
3236        assert_eq!(content, "");
3237
3238        // Only start line
3239        let content = thread
3240            .update(cx, |thread, cx| {
3241                thread.read_text_file(path!("/tmp/foo").into(), Some(1), None, false, cx)
3242            })
3243            .await
3244            .unwrap();
3245
3246        assert_eq!(content, "");
3247
3248        // Only limit
3249        let content = thread
3250            .update(cx, |thread, cx| {
3251                thread.read_text_file(path!("/tmp/foo").into(), None, Some(2), false, cx)
3252            })
3253            .await
3254            .unwrap();
3255
3256        assert_eq!(content, "");
3257
3258        // Range
3259        let content = thread
3260            .update(cx, |thread, cx| {
3261                thread.read_text_file(path!("/tmp/foo").into(), Some(1), Some(1), false, cx)
3262            })
3263            .await
3264            .unwrap();
3265
3266        assert_eq!(content, "");
3267
3268        // Invalid
3269        let err = thread
3270            .update(cx, |thread, cx| {
3271                thread.read_text_file(path!("/tmp/foo").into(), Some(5), Some(2), false, cx)
3272            })
3273            .await
3274            .unwrap_err();
3275
3276        assert_eq!(
3277            err.to_string(),
3278            "Invalid params: \"Attempting to read beyond the end of the file, line 1:0\""
3279        );
3280    }
3281    #[gpui::test]
3282    async fn test_reading_non_existing_file(cx: &mut TestAppContext) {
3283        init_test(cx);
3284
3285        let fs = FakeFs::new(cx.executor());
3286        fs.insert_tree(path!("/tmp"), json!({})).await;
3287        let project = Project::test(fs.clone(), [], cx).await;
3288        project
3289            .update(cx, |project, cx| {
3290                project.find_or_create_worktree(path!("/tmp"), true, cx)
3291            })
3292            .await
3293            .unwrap();
3294
3295        let connection = Rc::new(FakeAgentConnection::new());
3296
3297        let thread = cx
3298            .update(|cx| connection.new_session(project, Path::new(path!("/tmp")), cx))
3299            .await
3300            .unwrap();
3301
3302        // Out of project file
3303        let err = thread
3304            .update(cx, |thread, cx| {
3305                thread.read_text_file(path!("/foo").into(), None, None, false, cx)
3306            })
3307            .await
3308            .unwrap_err();
3309
3310        assert_eq!(err.code, acp::ErrorCode::ResourceNotFound);
3311    }
3312
3313    #[gpui::test]
3314    async fn test_succeeding_canceled_toolcall(cx: &mut TestAppContext) {
3315        init_test(cx);
3316
3317        let fs = FakeFs::new(cx.executor());
3318        let project = Project::test(fs, [], cx).await;
3319        let id = acp::ToolCallId::new("test");
3320
3321        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3322            let id = id.clone();
3323            move |_, thread, mut cx| {
3324                let id = id.clone();
3325                async move {
3326                    thread
3327                        .update(&mut cx, |thread, cx| {
3328                            thread.handle_session_update(
3329                                acp::SessionUpdate::ToolCall(
3330                                    acp::ToolCall::new(id.clone(), "Label")
3331                                        .kind(acp::ToolKind::Fetch)
3332                                        .status(acp::ToolCallStatus::InProgress),
3333                                ),
3334                                cx,
3335                            )
3336                        })
3337                        .unwrap()
3338                        .unwrap();
3339                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3340                }
3341                .boxed_local()
3342            }
3343        }));
3344
3345        let thread = cx
3346            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3347            .await
3348            .unwrap();
3349
3350        let request = thread.update(cx, |thread, cx| {
3351            thread.send_raw("Fetch https://example.com", cx)
3352        });
3353
3354        run_until_first_tool_call(&thread, cx).await;
3355
3356        thread.read_with(cx, |thread, _| {
3357            assert!(matches!(
3358                thread.entries[1],
3359                AgentThreadEntry::ToolCall(ToolCall {
3360                    status: ToolCallStatus::InProgress,
3361                    ..
3362                })
3363            ));
3364        });
3365
3366        thread.update(cx, |thread, cx| thread.cancel(cx)).await;
3367
3368        thread.read_with(cx, |thread, _| {
3369            assert!(matches!(
3370                &thread.entries[1],
3371                AgentThreadEntry::ToolCall(ToolCall {
3372                    status: ToolCallStatus::Canceled,
3373                    ..
3374                })
3375            ));
3376        });
3377
3378        thread
3379            .update(cx, |thread, cx| {
3380                thread.handle_session_update(
3381                    acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate::new(
3382                        id,
3383                        acp::ToolCallUpdateFields::new().status(acp::ToolCallStatus::Completed),
3384                    )),
3385                    cx,
3386                )
3387            })
3388            .unwrap();
3389
3390        request.await.unwrap();
3391
3392        thread.read_with(cx, |thread, _| {
3393            assert!(matches!(
3394                thread.entries[1],
3395                AgentThreadEntry::ToolCall(ToolCall {
3396                    status: ToolCallStatus::Completed,
3397                    ..
3398                })
3399            ));
3400        });
3401    }
3402
3403    #[gpui::test]
3404    async fn test_no_pending_edits_if_tool_calls_are_completed(cx: &mut TestAppContext) {
3405        init_test(cx);
3406        let fs = FakeFs::new(cx.background_executor.clone());
3407        fs.insert_tree(path!("/test"), json!({})).await;
3408        let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
3409
3410        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3411            move |_, thread, mut cx| {
3412                async move {
3413                    thread
3414                        .update(&mut cx, |thread, cx| {
3415                            thread.handle_session_update(
3416                                acp::SessionUpdate::ToolCall(
3417                                    acp::ToolCall::new("test", "Label")
3418                                        .kind(acp::ToolKind::Edit)
3419                                        .status(acp::ToolCallStatus::Completed)
3420                                        .content(vec![acp::ToolCallContent::Diff(acp::Diff::new(
3421                                            "/test/test.txt",
3422                                            "foo",
3423                                        ))]),
3424                                ),
3425                                cx,
3426                            )
3427                        })
3428                        .unwrap()
3429                        .unwrap();
3430                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3431                }
3432                .boxed_local()
3433            }
3434        }));
3435
3436        let thread = cx
3437            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3438            .await
3439            .unwrap();
3440
3441        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Hi".into()], cx)))
3442            .await
3443            .unwrap();
3444
3445        assert!(cx.read(|cx| !thread.read(cx).has_pending_edit_tool_calls()));
3446    }
3447
3448    #[gpui::test(iterations = 10)]
3449    async fn test_checkpoints(cx: &mut TestAppContext) {
3450        init_test(cx);
3451        let fs = FakeFs::new(cx.background_executor.clone());
3452        fs.insert_tree(
3453            path!("/test"),
3454            json!({
3455                ".git": {}
3456            }),
3457        )
3458        .await;
3459        let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
3460
3461        let simulate_changes = Arc::new(AtomicBool::new(true));
3462        let next_filename = Arc::new(AtomicUsize::new(0));
3463        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3464            let simulate_changes = simulate_changes.clone();
3465            let next_filename = next_filename.clone();
3466            let fs = fs.clone();
3467            move |request, thread, mut cx| {
3468                let fs = fs.clone();
3469                let simulate_changes = simulate_changes.clone();
3470                let next_filename = next_filename.clone();
3471                async move {
3472                    if simulate_changes.load(SeqCst) {
3473                        let filename = format!("/test/file-{}", next_filename.fetch_add(1, SeqCst));
3474                        fs.write(Path::new(&filename), b"").await?;
3475                    }
3476
3477                    let acp::ContentBlock::Text(content) = &request.prompt[0] else {
3478                        panic!("expected text content block");
3479                    };
3480                    thread.update(&mut cx, |thread, cx| {
3481                        thread
3482                            .handle_session_update(
3483                                acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(
3484                                    content.text.to_uppercase().into(),
3485                                )),
3486                                cx,
3487                            )
3488                            .unwrap();
3489                    })?;
3490                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3491                }
3492                .boxed_local()
3493            }
3494        }));
3495        let thread = cx
3496            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3497            .await
3498            .unwrap();
3499
3500        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["Lorem".into()], cx)))
3501            .await
3502            .unwrap();
3503        thread.read_with(cx, |thread, cx| {
3504            assert_eq!(
3505                thread.to_markdown(cx),
3506                indoc! {"
3507                    ## User (checkpoint)
3508
3509                    Lorem
3510
3511                    ## Assistant
3512
3513                    LOREM
3514
3515                "}
3516            );
3517        });
3518        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
3519
3520        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["ipsum".into()], cx)))
3521            .await
3522            .unwrap();
3523        thread.read_with(cx, |thread, cx| {
3524            assert_eq!(
3525                thread.to_markdown(cx),
3526                indoc! {"
3527                    ## User (checkpoint)
3528
3529                    Lorem
3530
3531                    ## Assistant
3532
3533                    LOREM
3534
3535                    ## User (checkpoint)
3536
3537                    ipsum
3538
3539                    ## Assistant
3540
3541                    IPSUM
3542
3543                "}
3544            );
3545        });
3546        assert_eq!(
3547            fs.files(),
3548            vec![
3549                Path::new(path!("/test/file-0")),
3550                Path::new(path!("/test/file-1"))
3551            ]
3552        );
3553
3554        // Checkpoint isn't stored when there are no changes.
3555        simulate_changes.store(false, SeqCst);
3556        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["dolor".into()], cx)))
3557            .await
3558            .unwrap();
3559        thread.read_with(cx, |thread, cx| {
3560            assert_eq!(
3561                thread.to_markdown(cx),
3562                indoc! {"
3563                    ## User (checkpoint)
3564
3565                    Lorem
3566
3567                    ## Assistant
3568
3569                    LOREM
3570
3571                    ## User (checkpoint)
3572
3573                    ipsum
3574
3575                    ## Assistant
3576
3577                    IPSUM
3578
3579                    ## User
3580
3581                    dolor
3582
3583                    ## Assistant
3584
3585                    DOLOR
3586
3587                "}
3588            );
3589        });
3590        assert_eq!(
3591            fs.files(),
3592            vec![
3593                Path::new(path!("/test/file-0")),
3594                Path::new(path!("/test/file-1"))
3595            ]
3596        );
3597
3598        // Rewinding the conversation truncates the history and restores the checkpoint.
3599        thread
3600            .update(cx, |thread, cx| {
3601                let AgentThreadEntry::UserMessage(message) = &thread.entries[2] else {
3602                    panic!("unexpected entries {:?}", thread.entries)
3603                };
3604                thread.restore_checkpoint(message.id.clone().unwrap(), cx)
3605            })
3606            .await
3607            .unwrap();
3608        thread.read_with(cx, |thread, cx| {
3609            assert_eq!(
3610                thread.to_markdown(cx),
3611                indoc! {"
3612                    ## User (checkpoint)
3613
3614                    Lorem
3615
3616                    ## Assistant
3617
3618                    LOREM
3619
3620                "}
3621            );
3622        });
3623        assert_eq!(fs.files(), vec![Path::new(path!("/test/file-0"))]);
3624    }
3625
3626    #[gpui::test]
3627    async fn test_tool_result_refusal(cx: &mut TestAppContext) {
3628        use std::sync::atomic::AtomicUsize;
3629        init_test(cx);
3630
3631        let fs = FakeFs::new(cx.executor());
3632        let project = Project::test(fs, None, cx).await;
3633
3634        // Create a connection that simulates refusal after tool result
3635        let prompt_count = Arc::new(AtomicUsize::new(0));
3636        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3637            let prompt_count = prompt_count.clone();
3638            move |_request, thread, mut cx| {
3639                let count = prompt_count.fetch_add(1, SeqCst);
3640                async move {
3641                    if count == 0 {
3642                        // First prompt: Generate a tool call with result
3643                        thread.update(&mut cx, |thread, cx| {
3644                            thread
3645                                .handle_session_update(
3646                                    acp::SessionUpdate::ToolCall(
3647                                        acp::ToolCall::new("tool1", "Test Tool")
3648                                            .kind(acp::ToolKind::Fetch)
3649                                            .status(acp::ToolCallStatus::Completed)
3650                                            .raw_input(serde_json::json!({"query": "test"}))
3651                                            .raw_output(serde_json::json!({"result": "inappropriate content"})),
3652                                    ),
3653                                    cx,
3654                                )
3655                                .unwrap();
3656                        })?;
3657
3658                        // Now return refusal because of the tool result
3659                        Ok(acp::PromptResponse::new(acp::StopReason::Refusal))
3660                    } else {
3661                        Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3662                    }
3663                }
3664                .boxed_local()
3665            }
3666        }));
3667
3668        let thread = cx
3669            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3670            .await
3671            .unwrap();
3672
3673        // Track if we see a Refusal event
3674        let saw_refusal_event = Arc::new(std::sync::Mutex::new(false));
3675        let saw_refusal_event_captured = saw_refusal_event.clone();
3676        thread.update(cx, |_thread, cx| {
3677            cx.subscribe(
3678                &thread,
3679                move |_thread, _event_thread, event: &AcpThreadEvent, _cx| {
3680                    if matches!(event, AcpThreadEvent::Refusal) {
3681                        *saw_refusal_event_captured.lock().unwrap() = true;
3682                    }
3683                },
3684            )
3685            .detach();
3686        });
3687
3688        // Send a user message - this will trigger tool call and then refusal
3689        let send_task = thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
3690        cx.background_executor.spawn(send_task).detach();
3691        cx.run_until_parked();
3692
3693        // Verify that:
3694        // 1. A Refusal event WAS emitted (because it's a tool result refusal, not user prompt)
3695        // 2. The user message was NOT truncated
3696        assert!(
3697            *saw_refusal_event.lock().unwrap(),
3698            "Refusal event should be emitted for tool result refusals"
3699        );
3700
3701        thread.read_with(cx, |thread, _| {
3702            let entries = thread.entries();
3703            assert!(entries.len() >= 2, "Should have user message and tool call");
3704
3705            // Verify user message is still there
3706            assert!(
3707                matches!(entries[0], AgentThreadEntry::UserMessage(_)),
3708                "User message should not be truncated"
3709            );
3710
3711            // Verify tool call is there with result
3712            if let AgentThreadEntry::ToolCall(tool_call) = &entries[1] {
3713                assert!(
3714                    tool_call.raw_output.is_some(),
3715                    "Tool call should have output"
3716                );
3717            } else {
3718                panic!("Expected tool call at index 1");
3719            }
3720        });
3721    }
3722
3723    #[gpui::test]
3724    async fn test_user_prompt_refusal_emits_event(cx: &mut TestAppContext) {
3725        init_test(cx);
3726
3727        let fs = FakeFs::new(cx.executor());
3728        let project = Project::test(fs, None, cx).await;
3729
3730        let refuse_next = Arc::new(AtomicBool::new(false));
3731        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3732            let refuse_next = refuse_next.clone();
3733            move |_request, _thread, _cx| {
3734                if refuse_next.load(SeqCst) {
3735                    async move { Ok(acp::PromptResponse::new(acp::StopReason::Refusal)) }
3736                        .boxed_local()
3737                } else {
3738                    async move { Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)) }
3739                        .boxed_local()
3740                }
3741            }
3742        }));
3743
3744        let thread = cx
3745            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3746            .await
3747            .unwrap();
3748
3749        // Track if we see a Refusal event
3750        let saw_refusal_event = Arc::new(std::sync::Mutex::new(false));
3751        let saw_refusal_event_captured = saw_refusal_event.clone();
3752        thread.update(cx, |_thread, cx| {
3753            cx.subscribe(
3754                &thread,
3755                move |_thread, _event_thread, event: &AcpThreadEvent, _cx| {
3756                    if matches!(event, AcpThreadEvent::Refusal) {
3757                        *saw_refusal_event_captured.lock().unwrap() = true;
3758                    }
3759                },
3760            )
3761            .detach();
3762        });
3763
3764        // Send a message that will be refused
3765        refuse_next.store(true, SeqCst);
3766        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx)))
3767            .await
3768            .unwrap();
3769
3770        // Verify that a Refusal event WAS emitted for user prompt refusal
3771        assert!(
3772            *saw_refusal_event.lock().unwrap(),
3773            "Refusal event should be emitted for user prompt refusals"
3774        );
3775
3776        // Verify the message was truncated (user prompt refusal)
3777        thread.read_with(cx, |thread, cx| {
3778            assert_eq!(thread.to_markdown(cx), "");
3779        });
3780    }
3781
3782    #[gpui::test]
3783    async fn test_refusal(cx: &mut TestAppContext) {
3784        init_test(cx);
3785        let fs = FakeFs::new(cx.background_executor.clone());
3786        fs.insert_tree(path!("/"), json!({})).await;
3787        let project = Project::test(fs.clone(), [path!("/").as_ref()], cx).await;
3788
3789        let refuse_next = Arc::new(AtomicBool::new(false));
3790        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
3791            let refuse_next = refuse_next.clone();
3792            move |request, thread, mut cx| {
3793                let refuse_next = refuse_next.clone();
3794                async move {
3795                    if refuse_next.load(SeqCst) {
3796                        return Ok(acp::PromptResponse::new(acp::StopReason::Refusal));
3797                    }
3798
3799                    let acp::ContentBlock::Text(content) = &request.prompt[0] else {
3800                        panic!("expected text content block");
3801                    };
3802                    thread.update(&mut cx, |thread, cx| {
3803                        thread
3804                            .handle_session_update(
3805                                acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(
3806                                    content.text.to_uppercase().into(),
3807                                )),
3808                                cx,
3809                            )
3810                            .unwrap();
3811                    })?;
3812                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
3813                }
3814                .boxed_local()
3815            }
3816        }));
3817        let thread = cx
3818            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
3819            .await
3820            .unwrap();
3821
3822        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["hello".into()], cx)))
3823            .await
3824            .unwrap();
3825        thread.read_with(cx, |thread, cx| {
3826            assert_eq!(
3827                thread.to_markdown(cx),
3828                indoc! {"
3829                    ## User
3830
3831                    hello
3832
3833                    ## Assistant
3834
3835                    HELLO
3836
3837                "}
3838            );
3839        });
3840
3841        // Simulate refusing the second message. The message should be truncated
3842        // when a user prompt is refused.
3843        refuse_next.store(true, SeqCst);
3844        cx.update(|cx| thread.update(cx, |thread, cx| thread.send(vec!["world".into()], cx)))
3845            .await
3846            .unwrap();
3847        thread.read_with(cx, |thread, cx| {
3848            assert_eq!(
3849                thread.to_markdown(cx),
3850                indoc! {"
3851                    ## User
3852
3853                    hello
3854
3855                    ## Assistant
3856
3857                    HELLO
3858
3859                "}
3860            );
3861        });
3862    }
3863
3864    async fn run_until_first_tool_call(
3865        thread: &Entity<AcpThread>,
3866        cx: &mut TestAppContext,
3867    ) -> usize {
3868        let (mut tx, mut rx) = mpsc::channel::<usize>(1);
3869
3870        let subscription = cx.update(|cx| {
3871            cx.subscribe(thread, move |thread, _, cx| {
3872                for (ix, entry) in thread.read(cx).entries.iter().enumerate() {
3873                    if matches!(entry, AgentThreadEntry::ToolCall(_)) {
3874                        return tx.try_send(ix).unwrap();
3875                    }
3876                }
3877            })
3878        });
3879
3880        select! {
3881            _ = futures::FutureExt::fuse(cx.background_executor.timer(Duration::from_secs(10))) => {
3882                panic!("Timeout waiting for tool call")
3883            }
3884            ix = rx.next().fuse() => {
3885                drop(subscription);
3886                ix.unwrap()
3887            }
3888        }
3889    }
3890
3891    #[derive(Clone, Default)]
3892    struct FakeAgentConnection {
3893        auth_methods: Vec<acp::AuthMethod>,
3894        sessions: Arc<parking_lot::Mutex<HashMap<acp::SessionId, WeakEntity<AcpThread>>>>,
3895        on_user_message: Option<
3896            Rc<
3897                dyn Fn(
3898                        acp::PromptRequest,
3899                        WeakEntity<AcpThread>,
3900                        AsyncApp,
3901                    ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
3902                    + 'static,
3903            >,
3904        >,
3905    }
3906
3907    impl FakeAgentConnection {
3908        fn new() -> Self {
3909            Self {
3910                auth_methods: Vec::new(),
3911                on_user_message: None,
3912                sessions: Arc::default(),
3913            }
3914        }
3915
3916        #[expect(unused)]
3917        fn with_auth_methods(mut self, auth_methods: Vec<acp::AuthMethod>) -> Self {
3918            self.auth_methods = auth_methods;
3919            self
3920        }
3921
3922        fn on_user_message(
3923            mut self,
3924            handler: impl Fn(
3925                acp::PromptRequest,
3926                WeakEntity<AcpThread>,
3927                AsyncApp,
3928            ) -> LocalBoxFuture<'static, Result<acp::PromptResponse>>
3929            + 'static,
3930        ) -> Self {
3931            self.on_user_message.replace(Rc::new(handler));
3932            self
3933        }
3934    }
3935
3936    impl AgentConnection for FakeAgentConnection {
3937        fn telemetry_id(&self) -> SharedString {
3938            "fake".into()
3939        }
3940
3941        fn auth_methods(&self) -> &[acp::AuthMethod] {
3942            &self.auth_methods
3943        }
3944
3945        fn new_session(
3946            self: Rc<Self>,
3947            project: Entity<Project>,
3948            _cwd: &Path,
3949            cx: &mut App,
3950        ) -> Task<gpui::Result<Entity<AcpThread>>> {
3951            let session_id = acp::SessionId::new(
3952                rand::rng()
3953                    .sample_iter(&distr::Alphanumeric)
3954                    .take(7)
3955                    .map(char::from)
3956                    .collect::<String>(),
3957            );
3958            let action_log = cx.new(|_| ActionLog::new(project.clone()));
3959            let thread = cx.new(|cx| {
3960                AcpThread::new(
3961                    None,
3962                    "Test",
3963                    self.clone(),
3964                    project,
3965                    action_log,
3966                    session_id.clone(),
3967                    watch::Receiver::constant(
3968                        acp::PromptCapabilities::new()
3969                            .image(true)
3970                            .audio(true)
3971                            .embedded_context(true),
3972                    ),
3973                    cx,
3974                )
3975            });
3976            self.sessions.lock().insert(session_id, thread.downgrade());
3977            Task::ready(Ok(thread))
3978        }
3979
3980        fn authenticate(&self, method: acp::AuthMethodId, _cx: &mut App) -> Task<gpui::Result<()>> {
3981            if self.auth_methods().iter().any(|m| m.id == method) {
3982                Task::ready(Ok(()))
3983            } else {
3984                Task::ready(Err(anyhow!("Invalid Auth Method")))
3985            }
3986        }
3987
3988        fn prompt(
3989            &self,
3990            _id: Option<UserMessageId>,
3991            params: acp::PromptRequest,
3992            cx: &mut App,
3993        ) -> Task<gpui::Result<acp::PromptResponse>> {
3994            let sessions = self.sessions.lock();
3995            let thread = sessions.get(&params.session_id).unwrap();
3996            if let Some(handler) = &self.on_user_message {
3997                let handler = handler.clone();
3998                let thread = thread.clone();
3999                cx.spawn(async move |cx| handler(params, thread, cx.clone()).await)
4000            } else {
4001                Task::ready(Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)))
4002            }
4003        }
4004
4005        fn cancel(&self, _session_id: &acp::SessionId, _cx: &mut App) {}
4006
4007        fn truncate(
4008            &self,
4009            session_id: &acp::SessionId,
4010            _cx: &App,
4011        ) -> Option<Rc<dyn AgentSessionTruncate>> {
4012            Some(Rc::new(FakeAgentSessionEditor {
4013                _session_id: session_id.clone(),
4014            }))
4015        }
4016
4017        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
4018            self
4019        }
4020    }
4021
4022    struct FakeAgentSessionEditor {
4023        _session_id: acp::SessionId,
4024    }
4025
4026    impl AgentSessionTruncate for FakeAgentSessionEditor {
4027        fn run(&self, _message_id: UserMessageId, _cx: &mut App) -> Task<Result<()>> {
4028            Task::ready(Ok(()))
4029        }
4030    }
4031
4032    #[gpui::test]
4033    async fn test_tool_call_not_found_creates_failed_entry(cx: &mut TestAppContext) {
4034        init_test(cx);
4035
4036        let fs = FakeFs::new(cx.executor());
4037        let project = Project::test(fs, [], cx).await;
4038        let connection = Rc::new(FakeAgentConnection::new());
4039        let thread = cx
4040            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
4041            .await
4042            .unwrap();
4043
4044        // Try to update a tool call that doesn't exist
4045        let nonexistent_id = acp::ToolCallId::new("nonexistent-tool-call");
4046        thread.update(cx, |thread, cx| {
4047            let result = thread.handle_session_update(
4048                acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate::new(
4049                    nonexistent_id.clone(),
4050                    acp::ToolCallUpdateFields::new().status(acp::ToolCallStatus::Completed),
4051                )),
4052                cx,
4053            );
4054
4055            // The update should succeed (not return an error)
4056            assert!(result.is_ok());
4057
4058            // There should now be exactly one entry in the thread
4059            assert_eq!(thread.entries.len(), 1);
4060
4061            // The entry should be a failed tool call
4062            if let AgentThreadEntry::ToolCall(tool_call) = &thread.entries[0] {
4063                assert_eq!(tool_call.id, nonexistent_id);
4064                assert!(matches!(tool_call.status, ToolCallStatus::Failed));
4065                assert_eq!(tool_call.kind, acp::ToolKind::Fetch);
4066
4067                // Check that the content contains the error message
4068                assert_eq!(tool_call.content.len(), 1);
4069                if let ToolCallContent::ContentBlock(content_block) = &tool_call.content[0] {
4070                    match content_block {
4071                        ContentBlock::Markdown { markdown } => {
4072                            let markdown_text = markdown.read(cx).source();
4073                            assert!(markdown_text.contains("Tool call not found"));
4074                        }
4075                        ContentBlock::Empty => panic!("Expected markdown content, got empty"),
4076                        ContentBlock::ResourceLink { .. } => {
4077                            panic!("Expected markdown content, got resource link")
4078                        }
4079                        ContentBlock::Image { .. } => {
4080                            panic!("Expected markdown content, got image")
4081                        }
4082                    }
4083                } else {
4084                    panic!("Expected ContentBlock, got: {:?}", tool_call.content[0]);
4085                }
4086            } else {
4087                panic!("Expected ToolCall entry, got: {:?}", thread.entries[0]);
4088            }
4089        });
4090    }
4091
4092    /// Tests that restoring a checkpoint properly cleans up terminals that were
4093    /// created after that checkpoint, and cancels any in-progress generation.
4094    ///
4095    /// Reproduces issue #35142: When a checkpoint is restored, any terminal processes
4096    /// that were started after that checkpoint should be terminated, and any in-progress
4097    /// AI generation should be canceled.
4098    #[gpui::test]
4099    async fn test_restore_checkpoint_kills_terminal(cx: &mut TestAppContext) {
4100        init_test(cx);
4101
4102        let fs = FakeFs::new(cx.executor());
4103        let project = Project::test(fs, [], cx).await;
4104        let connection = Rc::new(FakeAgentConnection::new());
4105        let thread = cx
4106            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
4107            .await
4108            .unwrap();
4109
4110        // Send first user message to create a checkpoint
4111        cx.update(|cx| {
4112            thread.update(cx, |thread, cx| {
4113                thread.send(vec!["first message".into()], cx)
4114            })
4115        })
4116        .await
4117        .unwrap();
4118
4119        // Send second message (creates another checkpoint) - we'll restore to this one
4120        cx.update(|cx| {
4121            thread.update(cx, |thread, cx| {
4122                thread.send(vec!["second message".into()], cx)
4123            })
4124        })
4125        .await
4126        .unwrap();
4127
4128        // Create 2 terminals BEFORE the checkpoint that have completed running
4129        let terminal_id_1 = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
4130        let mock_terminal_1 = cx.new(|cx| {
4131            let builder = ::terminal::TerminalBuilder::new_display_only(
4132                ::terminal::terminal_settings::CursorShape::default(),
4133                ::terminal::terminal_settings::AlternateScroll::On,
4134                None,
4135                0,
4136                cx.background_executor(),
4137                PathStyle::local(),
4138            )
4139            .unwrap();
4140            builder.subscribe(cx)
4141        });
4142
4143        thread.update(cx, |thread, cx| {
4144            thread.on_terminal_provider_event(
4145                TerminalProviderEvent::Created {
4146                    terminal_id: terminal_id_1.clone(),
4147                    label: "echo 'first'".to_string(),
4148                    cwd: Some(PathBuf::from("/test")),
4149                    output_byte_limit: None,
4150                    terminal: mock_terminal_1.clone(),
4151                },
4152                cx,
4153            );
4154        });
4155
4156        thread.update(cx, |thread, cx| {
4157            thread.on_terminal_provider_event(
4158                TerminalProviderEvent::Output {
4159                    terminal_id: terminal_id_1.clone(),
4160                    data: b"first\n".to_vec(),
4161                },
4162                cx,
4163            );
4164        });
4165
4166        thread.update(cx, |thread, cx| {
4167            thread.on_terminal_provider_event(
4168                TerminalProviderEvent::Exit {
4169                    terminal_id: terminal_id_1.clone(),
4170                    status: acp::TerminalExitStatus::new().exit_code(0),
4171                },
4172                cx,
4173            );
4174        });
4175
4176        let terminal_id_2 = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
4177        let mock_terminal_2 = cx.new(|cx| {
4178            let builder = ::terminal::TerminalBuilder::new_display_only(
4179                ::terminal::terminal_settings::CursorShape::default(),
4180                ::terminal::terminal_settings::AlternateScroll::On,
4181                None,
4182                0,
4183                cx.background_executor(),
4184                PathStyle::local(),
4185            )
4186            .unwrap();
4187            builder.subscribe(cx)
4188        });
4189
4190        thread.update(cx, |thread, cx| {
4191            thread.on_terminal_provider_event(
4192                TerminalProviderEvent::Created {
4193                    terminal_id: terminal_id_2.clone(),
4194                    label: "echo 'second'".to_string(),
4195                    cwd: Some(PathBuf::from("/test")),
4196                    output_byte_limit: None,
4197                    terminal: mock_terminal_2.clone(),
4198                },
4199                cx,
4200            );
4201        });
4202
4203        thread.update(cx, |thread, cx| {
4204            thread.on_terminal_provider_event(
4205                TerminalProviderEvent::Output {
4206                    terminal_id: terminal_id_2.clone(),
4207                    data: b"second\n".to_vec(),
4208                },
4209                cx,
4210            );
4211        });
4212
4213        thread.update(cx, |thread, cx| {
4214            thread.on_terminal_provider_event(
4215                TerminalProviderEvent::Exit {
4216                    terminal_id: terminal_id_2.clone(),
4217                    status: acp::TerminalExitStatus::new().exit_code(0),
4218                },
4219                cx,
4220            );
4221        });
4222
4223        // Get the second message ID to restore to
4224        let second_message_id = thread.read_with(cx, |thread, _| {
4225            // At this point we have:
4226            // - Index 0: First user message (with checkpoint)
4227            // - Index 1: Second user message (with checkpoint)
4228            // No assistant responses because FakeAgentConnection just returns EndTurn
4229            let AgentThreadEntry::UserMessage(message) = &thread.entries[1] else {
4230                panic!("expected user message at index 1");
4231            };
4232            message.id.clone().unwrap()
4233        });
4234
4235        // Create a terminal AFTER the checkpoint we'll restore to.
4236        // This simulates the AI agent starting a long-running terminal command.
4237        let terminal_id = acp::TerminalId::new(uuid::Uuid::new_v4().to_string());
4238        let mock_terminal = cx.new(|cx| {
4239            let builder = ::terminal::TerminalBuilder::new_display_only(
4240                ::terminal::terminal_settings::CursorShape::default(),
4241                ::terminal::terminal_settings::AlternateScroll::On,
4242                None,
4243                0,
4244                cx.background_executor(),
4245                PathStyle::local(),
4246            )
4247            .unwrap();
4248            builder.subscribe(cx)
4249        });
4250
4251        // Register the terminal as created
4252        thread.update(cx, |thread, cx| {
4253            thread.on_terminal_provider_event(
4254                TerminalProviderEvent::Created {
4255                    terminal_id: terminal_id.clone(),
4256                    label: "sleep 1000".to_string(),
4257                    cwd: Some(PathBuf::from("/test")),
4258                    output_byte_limit: None,
4259                    terminal: mock_terminal.clone(),
4260                },
4261                cx,
4262            );
4263        });
4264
4265        // Simulate the terminal producing output (still running)
4266        thread.update(cx, |thread, cx| {
4267            thread.on_terminal_provider_event(
4268                TerminalProviderEvent::Output {
4269                    terminal_id: terminal_id.clone(),
4270                    data: b"terminal is running...\n".to_vec(),
4271                },
4272                cx,
4273            );
4274        });
4275
4276        // Create a tool call entry that references this terminal
4277        // This represents the agent requesting a terminal command
4278        thread.update(cx, |thread, cx| {
4279            thread
4280                .handle_session_update(
4281                    acp::SessionUpdate::ToolCall(
4282                        acp::ToolCall::new("terminal-tool-1", "Running command")
4283                            .kind(acp::ToolKind::Execute)
4284                            .status(acp::ToolCallStatus::InProgress)
4285                            .content(vec![acp::ToolCallContent::Terminal(acp::Terminal::new(
4286                                terminal_id.clone(),
4287                            ))])
4288                            .raw_input(serde_json::json!({"command": "sleep 1000", "cd": "/test"})),
4289                    ),
4290                    cx,
4291                )
4292                .unwrap();
4293        });
4294
4295        // Verify terminal exists and is in the thread
4296        let terminal_exists_before =
4297            thread.read_with(cx, |thread, _| thread.terminals.contains_key(&terminal_id));
4298        assert!(
4299            terminal_exists_before,
4300            "Terminal should exist before checkpoint restore"
4301        );
4302
4303        // Verify the terminal's underlying task is still running (not completed)
4304        let terminal_running_before = thread.read_with(cx, |thread, _cx| {
4305            let terminal_entity = thread.terminals.get(&terminal_id).unwrap();
4306            terminal_entity.read_with(cx, |term, _cx| {
4307                term.output().is_none() // output is None means it's still running
4308            })
4309        });
4310        assert!(
4311            terminal_running_before,
4312            "Terminal should be running before checkpoint restore"
4313        );
4314
4315        // Verify we have the expected entries before restore
4316        let entry_count_before = thread.read_with(cx, |thread, _| thread.entries.len());
4317        assert!(
4318            entry_count_before > 1,
4319            "Should have multiple entries before restore"
4320        );
4321
4322        // Restore the checkpoint to the second message.
4323        // This should:
4324        // 1. Cancel any in-progress generation (via the cancel() call)
4325        // 2. Remove the terminal that was created after that point
4326        thread
4327            .update(cx, |thread, cx| {
4328                thread.restore_checkpoint(second_message_id, cx)
4329            })
4330            .await
4331            .unwrap();
4332
4333        // Verify that no send_task is in progress after restore
4334        // (cancel() clears the send_task)
4335        let has_send_task_after = thread.read_with(cx, |thread, _| thread.running_turn.is_some());
4336        assert!(
4337            !has_send_task_after,
4338            "Should not have a send_task after restore (cancel should have cleared it)"
4339        );
4340
4341        // Verify the entries were truncated (restoring to index 1 truncates at 1, keeping only index 0)
4342        let entry_count = thread.read_with(cx, |thread, _| thread.entries.len());
4343        assert_eq!(
4344            entry_count, 1,
4345            "Should have 1 entry after restore (only the first user message)"
4346        );
4347
4348        // Verify the 2 completed terminals from before the checkpoint still exist
4349        let terminal_1_exists = thread.read_with(cx, |thread, _| {
4350            thread.terminals.contains_key(&terminal_id_1)
4351        });
4352        assert!(
4353            terminal_1_exists,
4354            "Terminal 1 (from before checkpoint) should still exist"
4355        );
4356
4357        let terminal_2_exists = thread.read_with(cx, |thread, _| {
4358            thread.terminals.contains_key(&terminal_id_2)
4359        });
4360        assert!(
4361            terminal_2_exists,
4362            "Terminal 2 (from before checkpoint) should still exist"
4363        );
4364
4365        // Verify they're still in completed state
4366        let terminal_1_completed = thread.read_with(cx, |thread, _cx| {
4367            let terminal_entity = thread.terminals.get(&terminal_id_1).unwrap();
4368            terminal_entity.read_with(cx, |term, _cx| term.output().is_some())
4369        });
4370        assert!(terminal_1_completed, "Terminal 1 should still be completed");
4371
4372        let terminal_2_completed = thread.read_with(cx, |thread, _cx| {
4373            let terminal_entity = thread.terminals.get(&terminal_id_2).unwrap();
4374            terminal_entity.read_with(cx, |term, _cx| term.output().is_some())
4375        });
4376        assert!(terminal_2_completed, "Terminal 2 should still be completed");
4377
4378        // Verify the running terminal (created after checkpoint) was removed
4379        let terminal_3_exists =
4380            thread.read_with(cx, |thread, _| thread.terminals.contains_key(&terminal_id));
4381        assert!(
4382            !terminal_3_exists,
4383            "Terminal 3 (created after checkpoint) should have been removed"
4384        );
4385
4386        // Verify total count is 2 (the two from before the checkpoint)
4387        let terminal_count = thread.read_with(cx, |thread, _| thread.terminals.len());
4388        assert_eq!(
4389            terminal_count, 2,
4390            "Should have exactly 2 terminals (the completed ones from before checkpoint)"
4391        );
4392    }
4393
4394    /// Tests that update_last_checkpoint correctly updates the original message's checkpoint
4395    /// even when a new user message is added while the async checkpoint comparison is in progress.
4396    ///
4397    /// This is a regression test for a bug where update_last_checkpoint would fail with
4398    /// "no checkpoint" if a new user message (without a checkpoint) was added between when
4399    /// update_last_checkpoint started and when its async closure ran.
4400    #[gpui::test]
4401    async fn test_update_last_checkpoint_with_new_message_added(cx: &mut TestAppContext) {
4402        init_test(cx);
4403
4404        let fs = FakeFs::new(cx.executor());
4405        fs.insert_tree(path!("/test"), json!({".git": {}, "file.txt": "content"}))
4406            .await;
4407        let project = Project::test(fs.clone(), [Path::new(path!("/test"))], cx).await;
4408
4409        let handler_done = Arc::new(AtomicBool::new(false));
4410        let handler_done_clone = handler_done.clone();
4411        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
4412            move |_, _thread, _cx| {
4413                handler_done_clone.store(true, SeqCst);
4414                async move { Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)) }.boxed_local()
4415            },
4416        ));
4417
4418        let thread = cx
4419            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
4420            .await
4421            .unwrap();
4422
4423        let send_future = thread.update(cx, |thread, cx| thread.send_raw("First message", cx));
4424        let send_task = cx.background_executor.spawn(send_future);
4425
4426        // Tick until handler completes, then a few more to let update_last_checkpoint start
4427        while !handler_done.load(SeqCst) {
4428            cx.executor().tick();
4429        }
4430        for _ in 0..5 {
4431            cx.executor().tick();
4432        }
4433
4434        thread.update(cx, |thread, cx| {
4435            thread.push_entry(
4436                AgentThreadEntry::UserMessage(UserMessage {
4437                    id: Some(UserMessageId::new()),
4438                    content: ContentBlock::Empty,
4439                    chunks: vec!["Injected message (no checkpoint)".into()],
4440                    checkpoint: None,
4441                    indented: false,
4442                }),
4443                cx,
4444            );
4445        });
4446
4447        cx.run_until_parked();
4448        let result = send_task.await;
4449
4450        assert!(
4451            result.is_ok(),
4452            "send should succeed even when new message added during update_last_checkpoint: {:?}",
4453            result.err()
4454        );
4455    }
4456
4457    /// Tests that when a follow-up message is sent during generation,
4458    /// the first turn completing does NOT clear `running_turn` because
4459    /// it now belongs to the second turn.
4460    #[gpui::test]
4461    async fn test_follow_up_message_during_generation_does_not_clear_turn(cx: &mut TestAppContext) {
4462        init_test(cx);
4463
4464        let fs = FakeFs::new(cx.executor());
4465        let project = Project::test(fs, [], cx).await;
4466
4467        // First handler waits for this signal before completing
4468        let (first_complete_tx, first_complete_rx) = futures::channel::oneshot::channel::<()>();
4469        let first_complete_rx = RefCell::new(Some(first_complete_rx));
4470
4471        let connection = Rc::new(FakeAgentConnection::new().on_user_message({
4472            move |params, _thread, _cx| {
4473                let first_complete_rx = first_complete_rx.borrow_mut().take();
4474                let is_first = params
4475                    .prompt
4476                    .iter()
4477                    .any(|c| matches!(c, acp::ContentBlock::Text(t) if t.text.contains("first")));
4478
4479                async move {
4480                    if is_first {
4481                        // First handler waits until signaled
4482                        if let Some(rx) = first_complete_rx {
4483                            rx.await.ok();
4484                        }
4485                    }
4486                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
4487                }
4488                .boxed_local()
4489            }
4490        }));
4491
4492        let thread = cx
4493            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
4494            .await
4495            .unwrap();
4496
4497        // Send first message (turn_id=1) - handler will block
4498        let first_request = thread.update(cx, |thread, cx| thread.send_raw("first", cx));
4499        assert_eq!(thread.read_with(cx, |t, _| t.turn_id), 1);
4500
4501        // Send second message (turn_id=2) while first is still blocked
4502        // This calls cancel() which takes turn 1's running_turn and sets turn 2's
4503        let second_request = thread.update(cx, |thread, cx| thread.send_raw("second", cx));
4504        assert_eq!(thread.read_with(cx, |t, _| t.turn_id), 2);
4505
4506        let running_turn_after_second_send =
4507            thread.read_with(cx, |thread, _| thread.running_turn.as_ref().map(|t| t.id));
4508        assert_eq!(
4509            running_turn_after_second_send,
4510            Some(2),
4511            "running_turn should be set to turn 2 after sending second message"
4512        );
4513
4514        // Now signal first handler to complete
4515        first_complete_tx.send(()).ok();
4516
4517        // First request completes - should NOT clear running_turn
4518        // because running_turn now belongs to turn 2
4519        first_request.await.unwrap();
4520
4521        let running_turn_after_first =
4522            thread.read_with(cx, |thread, _| thread.running_turn.as_ref().map(|t| t.id));
4523        assert_eq!(
4524            running_turn_after_first,
4525            Some(2),
4526            "first turn completing should not clear running_turn (belongs to turn 2)"
4527        );
4528
4529        // Second request completes - SHOULD clear running_turn
4530        second_request.await.unwrap();
4531
4532        let running_turn_after_second =
4533            thread.read_with(cx, |thread, _| thread.running_turn.is_some());
4534        assert!(
4535            !running_turn_after_second,
4536            "second turn completing should clear running_turn"
4537        );
4538    }
4539
4540    #[gpui::test]
4541    async fn test_send_returns_cancelled_response_and_marks_tools_as_cancelled(
4542        cx: &mut TestAppContext,
4543    ) {
4544        init_test(cx);
4545
4546        let fs = FakeFs::new(cx.executor());
4547        let project = Project::test(fs, [], cx).await;
4548
4549        let connection = Rc::new(FakeAgentConnection::new().on_user_message(
4550            move |_params, thread, mut cx| {
4551                async move {
4552                    thread
4553                        .update(&mut cx, |thread, cx| {
4554                            thread.handle_session_update(
4555                                acp::SessionUpdate::ToolCall(
4556                                    acp::ToolCall::new(
4557                                        acp::ToolCallId::new("test-tool"),
4558                                        "Test Tool",
4559                                    )
4560                                    .kind(acp::ToolKind::Fetch)
4561                                    .status(acp::ToolCallStatus::InProgress),
4562                                ),
4563                                cx,
4564                            )
4565                        })
4566                        .unwrap()
4567                        .unwrap();
4568
4569                    Ok(acp::PromptResponse::new(acp::StopReason::Cancelled))
4570                }
4571                .boxed_local()
4572            },
4573        ));
4574
4575        let thread = cx
4576            .update(|cx| connection.new_session(project, Path::new(path!("/test")), cx))
4577            .await
4578            .unwrap();
4579
4580        let response = thread
4581            .update(cx, |thread, cx| thread.send_raw("test message", cx))
4582            .await;
4583
4584        let response = response
4585            .expect("send should succeed")
4586            .expect("should have response");
4587        assert_eq!(
4588            response.stop_reason,
4589            acp::StopReason::Cancelled,
4590            "response should have Cancelled stop_reason"
4591        );
4592
4593        thread.read_with(cx, |thread, _| {
4594            let tool_entry = thread
4595                .entries
4596                .iter()
4597                .find_map(|e| {
4598                    if let AgentThreadEntry::ToolCall(call) = e {
4599                        Some(call)
4600                    } else {
4601                        None
4602                    }
4603                })
4604                .expect("should have tool call entry");
4605
4606            assert!(
4607                matches!(tool_entry.status, ToolCallStatus::Canceled),
4608                "tool should be marked as Canceled when response is Cancelled, got {:?}",
4609                tool_entry.status
4610            );
4611        });
4612    }
4613}