mod.rs

   1use super::*;
   2use acp_thread::{
   3    AgentConnection, AgentModelGroupName, AgentModelList, PermissionOptions, ThreadStatus,
   4    UserMessageId,
   5};
   6use agent_client_protocol::{self as acp};
   7use agent_settings::AgentProfileId;
   8use anyhow::Result;
   9use client::{Client, RefreshLlmTokenListener, UserStore};
  10use collections::IndexMap;
  11use context_server::{ContextServer, ContextServerCommand, ContextServerId};
  12
  13use fs::{FakeFs, Fs};
  14use futures::{
  15    FutureExt as _, StreamExt,
  16    channel::{
  17        mpsc::{self, UnboundedReceiver},
  18        oneshot,
  19    },
  20    future::{Fuse, Shared},
  21};
  22use gpui::{
  23    App, AppContext, AsyncApp, Entity, Task, TestAppContext, UpdateGlobal,
  24    http_client::FakeHttpClient,
  25};
  26use indoc::indoc;
  27use language_model::{
  28    CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
  29    LanguageModelId, LanguageModelProviderName, LanguageModelRegistry, LanguageModelRequest,
  30    LanguageModelRequestMessage, LanguageModelToolResult, LanguageModelToolSchemaFormat,
  31    LanguageModelToolUse, MessageContent, Role, StopReason, TokenUsage,
  32    fake_provider::FakeLanguageModel,
  33};
  34use pretty_assertions::assert_eq;
  35use project::{
  36    Project, context_server_store::ContextServerStore, project_settings::ProjectSettings,
  37};
  38use prompt_store::ProjectContext;
  39use reqwest_client::ReqwestClient;
  40use schemars::JsonSchema;
  41use serde::{Deserialize, Serialize};
  42use serde_json::json;
  43use settings::{Settings, SettingsStore};
  44use std::{
  45    path::Path,
  46    pin::Pin,
  47    rc::Rc,
  48    sync::{
  49        Arc,
  50        atomic::{AtomicBool, AtomicUsize, Ordering},
  51    },
  52    time::Duration,
  53};
  54use util::path;
  55
  56mod edit_file_thread_test;
  57mod test_tools;
  58use test_tools::*;
  59
  60pub(crate) fn init_test(cx: &mut TestAppContext) {
  61    cx.update(|cx| {
  62        let settings_store = SettingsStore::test(cx);
  63        cx.set_global(settings_store);
  64    });
  65}
  66
  67pub(crate) struct FakeTerminalHandle {
  68    killed: Arc<AtomicBool>,
  69    stopped_by_user: Arc<AtomicBool>,
  70    exit_sender: std::cell::RefCell<Option<futures::channel::oneshot::Sender<()>>>,
  71    wait_for_exit: Shared<Task<acp::TerminalExitStatus>>,
  72    output: acp::TerminalOutputResponse,
  73    id: acp::TerminalId,
  74}
  75
  76impl FakeTerminalHandle {
  77    pub(crate) fn new_never_exits(cx: &mut App) -> Self {
  78        let killed = Arc::new(AtomicBool::new(false));
  79        let stopped_by_user = Arc::new(AtomicBool::new(false));
  80
  81        let (exit_sender, exit_receiver) = futures::channel::oneshot::channel();
  82
  83        let wait_for_exit = cx
  84            .spawn(async move |_cx| {
  85                // Wait for the exit signal (sent when kill() is called)
  86                let _ = exit_receiver.await;
  87                acp::TerminalExitStatus::new()
  88            })
  89            .shared();
  90
  91        Self {
  92            killed,
  93            stopped_by_user,
  94            exit_sender: std::cell::RefCell::new(Some(exit_sender)),
  95            wait_for_exit,
  96            output: acp::TerminalOutputResponse::new("partial output".to_string(), false),
  97            id: acp::TerminalId::new("fake_terminal".to_string()),
  98        }
  99    }
 100
 101    pub(crate) fn new_with_immediate_exit(cx: &mut App, exit_code: u32) -> Self {
 102        let killed = Arc::new(AtomicBool::new(false));
 103        let stopped_by_user = Arc::new(AtomicBool::new(false));
 104        let (exit_sender, _exit_receiver) = futures::channel::oneshot::channel();
 105
 106        let wait_for_exit = cx
 107            .spawn(async move |_cx| acp::TerminalExitStatus::new().exit_code(exit_code))
 108            .shared();
 109
 110        Self {
 111            killed,
 112            stopped_by_user,
 113            exit_sender: std::cell::RefCell::new(Some(exit_sender)),
 114            wait_for_exit,
 115            output: acp::TerminalOutputResponse::new("command output".to_string(), false),
 116            id: acp::TerminalId::new("fake_terminal".to_string()),
 117        }
 118    }
 119
 120    pub(crate) fn was_killed(&self) -> bool {
 121        self.killed.load(Ordering::SeqCst)
 122    }
 123
 124    pub(crate) fn set_stopped_by_user(&self, stopped: bool) {
 125        self.stopped_by_user.store(stopped, Ordering::SeqCst);
 126    }
 127
 128    pub(crate) fn signal_exit(&self) {
 129        if let Some(sender) = self.exit_sender.borrow_mut().take() {
 130            let _ = sender.send(());
 131        }
 132    }
 133}
 134
 135impl crate::TerminalHandle for FakeTerminalHandle {
 136    fn id(&self, _cx: &AsyncApp) -> Result<acp::TerminalId> {
 137        Ok(self.id.clone())
 138    }
 139
 140    fn current_output(&self, _cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
 141        Ok(self.output.clone())
 142    }
 143
 144    fn wait_for_exit(&self, _cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
 145        Ok(self.wait_for_exit.clone())
 146    }
 147
 148    fn kill(&self, _cx: &AsyncApp) -> Result<()> {
 149        self.killed.store(true, Ordering::SeqCst);
 150        self.signal_exit();
 151        Ok(())
 152    }
 153
 154    fn was_stopped_by_user(&self, _cx: &AsyncApp) -> Result<bool> {
 155        Ok(self.stopped_by_user.load(Ordering::SeqCst))
 156    }
 157}
 158
 159struct FakeSubagentHandle {
 160    session_id: acp::SessionId,
 161    send_task: Shared<Task<String>>,
 162}
 163
 164impl SubagentHandle for FakeSubagentHandle {
 165    fn id(&self) -> acp::SessionId {
 166        self.session_id.clone()
 167    }
 168
 169    fn num_entries(&self, _cx: &App) -> usize {
 170        unimplemented!()
 171    }
 172
 173    fn send(&self, _message: String, cx: &AsyncApp) -> Task<Result<String>> {
 174        let task = self.send_task.clone();
 175        cx.background_spawn(async move { Ok(task.await) })
 176    }
 177}
 178
 179#[derive(Default)]
 180pub(crate) struct FakeThreadEnvironment {
 181    terminal_handle: Option<Rc<FakeTerminalHandle>>,
 182    subagent_handle: Option<Rc<FakeSubagentHandle>>,
 183    terminal_creations: Arc<AtomicUsize>,
 184}
 185
 186impl FakeThreadEnvironment {
 187    pub(crate) fn with_terminal(self, terminal_handle: FakeTerminalHandle) -> Self {
 188        Self {
 189            terminal_handle: Some(terminal_handle.into()),
 190            ..self
 191        }
 192    }
 193
 194    pub(crate) fn terminal_creation_count(&self) -> usize {
 195        self.terminal_creations.load(Ordering::SeqCst)
 196    }
 197}
 198
 199impl crate::ThreadEnvironment for FakeThreadEnvironment {
 200    fn create_terminal(
 201        &self,
 202        _command: String,
 203        _cwd: Option<std::path::PathBuf>,
 204        _output_byte_limit: Option<u64>,
 205        _cx: &mut AsyncApp,
 206    ) -> Task<Result<Rc<dyn crate::TerminalHandle>>> {
 207        self.terminal_creations.fetch_add(1, Ordering::SeqCst);
 208        let handle = self
 209            .terminal_handle
 210            .clone()
 211            .expect("Terminal handle not available on FakeThreadEnvironment");
 212        Task::ready(Ok(handle as Rc<dyn crate::TerminalHandle>))
 213    }
 214
 215    fn create_subagent(&self, _label: String, _cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {
 216        Ok(self
 217            .subagent_handle
 218            .clone()
 219            .expect("Subagent handle not available on FakeThreadEnvironment")
 220            as Rc<dyn SubagentHandle>)
 221    }
 222}
 223
 224/// Environment that creates multiple independent terminal handles for testing concurrent terminals.
 225struct MultiTerminalEnvironment {
 226    handles: std::cell::RefCell<Vec<Rc<FakeTerminalHandle>>>,
 227}
 228
 229impl MultiTerminalEnvironment {
 230    fn new() -> Self {
 231        Self {
 232            handles: std::cell::RefCell::new(Vec::new()),
 233        }
 234    }
 235
 236    fn handles(&self) -> Vec<Rc<FakeTerminalHandle>> {
 237        self.handles.borrow().clone()
 238    }
 239}
 240
 241impl crate::ThreadEnvironment for MultiTerminalEnvironment {
 242    fn create_terminal(
 243        &self,
 244        _command: String,
 245        _cwd: Option<std::path::PathBuf>,
 246        _output_byte_limit: Option<u64>,
 247        cx: &mut AsyncApp,
 248    ) -> Task<Result<Rc<dyn crate::TerminalHandle>>> {
 249        let handle = Rc::new(cx.update(|cx| FakeTerminalHandle::new_never_exits(cx)));
 250        self.handles.borrow_mut().push(handle.clone());
 251        Task::ready(Ok(handle as Rc<dyn crate::TerminalHandle>))
 252    }
 253
 254    fn create_subagent(&self, _label: String, _cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {
 255        unimplemented!()
 256    }
 257}
 258
 259fn always_allow_tools(cx: &mut TestAppContext) {
 260    cx.update(|cx| {
 261        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
 262        settings.tool_permissions.default = settings::ToolPermissionMode::Allow;
 263        agent_settings::AgentSettings::override_global(settings, cx);
 264    });
 265}
 266
 267#[gpui::test]
 268async fn test_echo(cx: &mut TestAppContext) {
 269    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 270    let fake_model = model.as_fake();
 271
 272    let events = thread
 273        .update(cx, |thread, cx| {
 274            thread.send(UserMessageId::new(), ["Testing: Reply with 'Hello'"], cx)
 275        })
 276        .unwrap();
 277    cx.run_until_parked();
 278    fake_model.send_last_completion_stream_text_chunk("Hello");
 279    fake_model
 280        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
 281    fake_model.end_last_completion_stream();
 282
 283    let events = events.collect().await;
 284    thread.update(cx, |thread, _cx| {
 285        assert_eq!(
 286            thread.last_received_or_pending_message().unwrap().role(),
 287            Role::Assistant
 288        );
 289        assert_eq!(
 290            thread
 291                .last_received_or_pending_message()
 292                .unwrap()
 293                .to_markdown(),
 294            "Hello\n"
 295        )
 296    });
 297    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
 298}
 299
 300#[gpui::test]
 301async fn test_terminal_tool_timeout_kills_handle(cx: &mut TestAppContext) {
 302    init_test(cx);
 303    always_allow_tools(cx);
 304
 305    let fs = FakeFs::new(cx.executor());
 306    let project = Project::test(fs, [], cx).await;
 307
 308    let environment = Rc::new(cx.update(|cx| {
 309        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
 310    }));
 311    let handle = environment.terminal_handle.clone().unwrap();
 312
 313    #[allow(clippy::arc_with_non_send_sync)]
 314    let tool = Arc::new(crate::TerminalTool::new(project, environment));
 315    let (event_stream, mut rx) = crate::ToolCallEventStream::test();
 316
 317    let task = cx.update(|cx| {
 318        tool.run(
 319            ToolInput::resolved(crate::TerminalToolInput {
 320                command: "sleep 1000".to_string(),
 321                cd: ".".to_string(),
 322                timeout_ms: Some(5),
 323            }),
 324            event_stream,
 325            cx,
 326        )
 327    });
 328
 329    let update = rx.expect_update_fields().await;
 330    assert!(
 331        update.content.iter().any(|blocks| {
 332            blocks
 333                .iter()
 334                .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))
 335        }),
 336        "expected tool call update to include terminal content"
 337    );
 338
 339    let mut task_future: Pin<Box<Fuse<Task<Result<String, String>>>>> = Box::pin(task.fuse());
 340
 341    let deadline = std::time::Instant::now() + Duration::from_millis(500);
 342    loop {
 343        if let Some(result) = task_future.as_mut().now_or_never() {
 344            let result = result.expect("terminal tool task should complete");
 345
 346            assert!(
 347                handle.was_killed(),
 348                "expected terminal handle to be killed on timeout"
 349            );
 350            assert!(
 351                result.contains("partial output"),
 352                "expected result to include terminal output, got: {result}"
 353            );
 354            return;
 355        }
 356
 357        if std::time::Instant::now() >= deadline {
 358            panic!("timed out waiting for terminal tool task to complete");
 359        }
 360
 361        cx.run_until_parked();
 362        cx.background_executor.timer(Duration::from_millis(1)).await;
 363    }
 364}
 365
 366#[gpui::test]
 367#[ignore]
 368async fn test_terminal_tool_without_timeout_does_not_kill_handle(cx: &mut TestAppContext) {
 369    init_test(cx);
 370    always_allow_tools(cx);
 371
 372    let fs = FakeFs::new(cx.executor());
 373    let project = Project::test(fs, [], cx).await;
 374
 375    let environment = Rc::new(cx.update(|cx| {
 376        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
 377    }));
 378    let handle = environment.terminal_handle.clone().unwrap();
 379
 380    #[allow(clippy::arc_with_non_send_sync)]
 381    let tool = Arc::new(crate::TerminalTool::new(project, environment));
 382    let (event_stream, mut rx) = crate::ToolCallEventStream::test();
 383
 384    let _task = cx.update(|cx| {
 385        tool.run(
 386            ToolInput::resolved(crate::TerminalToolInput {
 387                command: "sleep 1000".to_string(),
 388                cd: ".".to_string(),
 389                timeout_ms: None,
 390            }),
 391            event_stream,
 392            cx,
 393        )
 394    });
 395
 396    let update = rx.expect_update_fields().await;
 397    assert!(
 398        update.content.iter().any(|blocks| {
 399            blocks
 400                .iter()
 401                .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))
 402        }),
 403        "expected tool call update to include terminal content"
 404    );
 405
 406    cx.background_executor
 407        .timer(Duration::from_millis(25))
 408        .await;
 409
 410    assert!(
 411        !handle.was_killed(),
 412        "did not expect terminal handle to be killed without a timeout"
 413    );
 414}
 415
 416#[gpui::test]
 417async fn test_thinking(cx: &mut TestAppContext) {
 418    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 419    let fake_model = model.as_fake();
 420
 421    let events = thread
 422        .update(cx, |thread, cx| {
 423            thread.send(
 424                UserMessageId::new(),
 425                [indoc! {"
 426                    Testing:
 427
 428                    Generate a thinking step where you just think the word 'Think',
 429                    and have your final answer be 'Hello'
 430                "}],
 431                cx,
 432            )
 433        })
 434        .unwrap();
 435    cx.run_until_parked();
 436    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::Thinking {
 437        text: "Think".to_string(),
 438        signature: None,
 439    });
 440    fake_model.send_last_completion_stream_text_chunk("Hello");
 441    fake_model
 442        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
 443    fake_model.end_last_completion_stream();
 444
 445    let events = events.collect().await;
 446    thread.update(cx, |thread, _cx| {
 447        assert_eq!(
 448            thread.last_received_or_pending_message().unwrap().role(),
 449            Role::Assistant
 450        );
 451        assert_eq!(
 452            thread
 453                .last_received_or_pending_message()
 454                .unwrap()
 455                .to_markdown(),
 456            indoc! {"
 457                <think>Think</think>
 458                Hello
 459            "}
 460        )
 461    });
 462    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
 463}
 464
 465#[gpui::test]
 466async fn test_system_prompt(cx: &mut TestAppContext) {
 467    let ThreadTest {
 468        model,
 469        thread,
 470        project_context,
 471        ..
 472    } = setup(cx, TestModel::Fake).await;
 473    let fake_model = model.as_fake();
 474
 475    project_context.update(cx, |project_context, _cx| {
 476        project_context.shell = "test-shell".into()
 477    });
 478    thread.update(cx, |thread, _| thread.add_tool(EchoTool));
 479    thread
 480        .update(cx, |thread, cx| {
 481            thread.send(UserMessageId::new(), ["abc"], cx)
 482        })
 483        .unwrap();
 484    cx.run_until_parked();
 485    let mut pending_completions = fake_model.pending_completions();
 486    assert_eq!(
 487        pending_completions.len(),
 488        1,
 489        "unexpected pending completions: {:?}",
 490        pending_completions
 491    );
 492
 493    let pending_completion = pending_completions.pop().unwrap();
 494    assert_eq!(pending_completion.messages[0].role, Role::System);
 495
 496    let system_message = &pending_completion.messages[0];
 497    let system_prompt = system_message.content[0].to_str().unwrap();
 498    assert!(
 499        system_prompt.contains("test-shell"),
 500        "unexpected system message: {:?}",
 501        system_message
 502    );
 503    assert!(
 504        system_prompt.contains("## Fixing Diagnostics"),
 505        "unexpected system message: {:?}",
 506        system_message
 507    );
 508}
 509
 510#[gpui::test]
 511async fn test_system_prompt_without_tools(cx: &mut TestAppContext) {
 512    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 513    let fake_model = model.as_fake();
 514
 515    thread
 516        .update(cx, |thread, cx| {
 517            thread.send(UserMessageId::new(), ["abc"], cx)
 518        })
 519        .unwrap();
 520    cx.run_until_parked();
 521    let mut pending_completions = fake_model.pending_completions();
 522    assert_eq!(
 523        pending_completions.len(),
 524        1,
 525        "unexpected pending completions: {:?}",
 526        pending_completions
 527    );
 528
 529    let pending_completion = pending_completions.pop().unwrap();
 530    assert_eq!(pending_completion.messages[0].role, Role::System);
 531
 532    let system_message = &pending_completion.messages[0];
 533    let system_prompt = system_message.content[0].to_str().unwrap();
 534    assert!(
 535        !system_prompt.contains("## Tool Use"),
 536        "unexpected system message: {:?}",
 537        system_message
 538    );
 539    assert!(
 540        !system_prompt.contains("## Fixing Diagnostics"),
 541        "unexpected system message: {:?}",
 542        system_message
 543    );
 544}
 545
 546#[gpui::test]
 547async fn test_prompt_caching(cx: &mut TestAppContext) {
 548    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 549    let fake_model = model.as_fake();
 550
 551    // Send initial user message and verify it's cached
 552    thread
 553        .update(cx, |thread, cx| {
 554            thread.send(UserMessageId::new(), ["Message 1"], cx)
 555        })
 556        .unwrap();
 557    cx.run_until_parked();
 558
 559    let completion = fake_model.pending_completions().pop().unwrap();
 560    assert_eq!(
 561        completion.messages[1..],
 562        vec![LanguageModelRequestMessage {
 563            role: Role::User,
 564            content: vec!["Message 1".into()],
 565            cache: true,
 566            reasoning_details: None,
 567        }]
 568    );
 569    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::Text(
 570        "Response to Message 1".into(),
 571    ));
 572    fake_model.end_last_completion_stream();
 573    cx.run_until_parked();
 574
 575    // Send another user message and verify only the latest is cached
 576    thread
 577        .update(cx, |thread, cx| {
 578            thread.send(UserMessageId::new(), ["Message 2"], cx)
 579        })
 580        .unwrap();
 581    cx.run_until_parked();
 582
 583    let completion = fake_model.pending_completions().pop().unwrap();
 584    assert_eq!(
 585        completion.messages[1..],
 586        vec![
 587            LanguageModelRequestMessage {
 588                role: Role::User,
 589                content: vec!["Message 1".into()],
 590                cache: false,
 591                reasoning_details: None,
 592            },
 593            LanguageModelRequestMessage {
 594                role: Role::Assistant,
 595                content: vec!["Response to Message 1".into()],
 596                cache: false,
 597                reasoning_details: None,
 598            },
 599            LanguageModelRequestMessage {
 600                role: Role::User,
 601                content: vec!["Message 2".into()],
 602                cache: true,
 603                reasoning_details: None,
 604            }
 605        ]
 606    );
 607    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::Text(
 608        "Response to Message 2".into(),
 609    ));
 610    fake_model.end_last_completion_stream();
 611    cx.run_until_parked();
 612
 613    // Simulate a tool call and verify that the latest tool result is cached
 614    thread.update(cx, |thread, _| thread.add_tool(EchoTool));
 615    thread
 616        .update(cx, |thread, cx| {
 617            thread.send(UserMessageId::new(), ["Use the echo tool"], cx)
 618        })
 619        .unwrap();
 620    cx.run_until_parked();
 621
 622    let tool_use = LanguageModelToolUse {
 623        id: "tool_1".into(),
 624        name: EchoTool::NAME.into(),
 625        raw_input: json!({"text": "test"}).to_string(),
 626        input: json!({"text": "test"}),
 627        is_input_complete: true,
 628        thought_signature: None,
 629    };
 630    fake_model
 631        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(tool_use.clone()));
 632    fake_model.end_last_completion_stream();
 633    cx.run_until_parked();
 634
 635    let completion = fake_model.pending_completions().pop().unwrap();
 636    let tool_result = LanguageModelToolResult {
 637        tool_use_id: "tool_1".into(),
 638        tool_name: EchoTool::NAME.into(),
 639        is_error: false,
 640        content: "test".into(),
 641        output: Some("test".into()),
 642    };
 643    assert_eq!(
 644        completion.messages[1..],
 645        vec![
 646            LanguageModelRequestMessage {
 647                role: Role::User,
 648                content: vec!["Message 1".into()],
 649                cache: false,
 650                reasoning_details: None,
 651            },
 652            LanguageModelRequestMessage {
 653                role: Role::Assistant,
 654                content: vec!["Response to Message 1".into()],
 655                cache: false,
 656                reasoning_details: None,
 657            },
 658            LanguageModelRequestMessage {
 659                role: Role::User,
 660                content: vec!["Message 2".into()],
 661                cache: false,
 662                reasoning_details: None,
 663            },
 664            LanguageModelRequestMessage {
 665                role: Role::Assistant,
 666                content: vec!["Response to Message 2".into()],
 667                cache: false,
 668                reasoning_details: None,
 669            },
 670            LanguageModelRequestMessage {
 671                role: Role::User,
 672                content: vec!["Use the echo tool".into()],
 673                cache: false,
 674                reasoning_details: None,
 675            },
 676            LanguageModelRequestMessage {
 677                role: Role::Assistant,
 678                content: vec![MessageContent::ToolUse(tool_use)],
 679                cache: false,
 680                reasoning_details: None,
 681            },
 682            LanguageModelRequestMessage {
 683                role: Role::User,
 684                content: vec![MessageContent::ToolResult(tool_result)],
 685                cache: true,
 686                reasoning_details: None,
 687            }
 688        ]
 689    );
 690}
 691
 692#[gpui::test]
 693#[cfg_attr(not(feature = "e2e"), ignore)]
 694async fn test_basic_tool_calls(cx: &mut TestAppContext) {
 695    let ThreadTest { thread, .. } = setup(cx, TestModel::Sonnet4).await;
 696
 697    // Test a tool call that's likely to complete *before* streaming stops.
 698    let events = thread
 699        .update(cx, |thread, cx| {
 700            thread.add_tool(EchoTool);
 701            thread.send(
 702                UserMessageId::new(),
 703                ["Now test the echo tool with 'Hello'. Does it work? Say 'Yes' or 'No'."],
 704                cx,
 705            )
 706        })
 707        .unwrap()
 708        .collect()
 709        .await;
 710    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
 711
 712    // Test a tool calls that's likely to complete *after* streaming stops.
 713    let events = thread
 714        .update(cx, |thread, cx| {
 715            thread.remove_tool(&EchoTool::NAME);
 716            thread.add_tool(DelayTool);
 717            thread.send(
 718                UserMessageId::new(),
 719                [
 720                    "Now call the delay tool with 200ms.",
 721                    "When the timer goes off, then you echo the output of the tool.",
 722                ],
 723                cx,
 724            )
 725        })
 726        .unwrap()
 727        .collect()
 728        .await;
 729    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
 730    thread.update(cx, |thread, _cx| {
 731        assert!(
 732            thread
 733                .last_received_or_pending_message()
 734                .unwrap()
 735                .as_agent_message()
 736                .unwrap()
 737                .content
 738                .iter()
 739                .any(|content| {
 740                    if let AgentMessageContent::Text(text) = content {
 741                        text.contains("Ding")
 742                    } else {
 743                        false
 744                    }
 745                }),
 746            "{}",
 747            thread.to_markdown()
 748        );
 749    });
 750}
 751
 752#[gpui::test]
 753#[cfg_attr(not(feature = "e2e"), ignore)]
 754async fn test_streaming_tool_calls(cx: &mut TestAppContext) {
 755    let ThreadTest { thread, .. } = setup(cx, TestModel::Sonnet4).await;
 756
 757    // Test a tool call that's likely to complete *before* streaming stops.
 758    let mut events = thread
 759        .update(cx, |thread, cx| {
 760            thread.add_tool(WordListTool);
 761            thread.send(UserMessageId::new(), ["Test the word_list tool."], cx)
 762        })
 763        .unwrap();
 764
 765    let mut saw_partial_tool_use = false;
 766    while let Some(event) = events.next().await {
 767        if let Ok(ThreadEvent::ToolCall(tool_call)) = event {
 768            thread.update(cx, |thread, _cx| {
 769                // Look for a tool use in the thread's last message
 770                let message = thread.last_received_or_pending_message().unwrap();
 771                let agent_message = message.as_agent_message().unwrap();
 772                let last_content = agent_message.content.last().unwrap();
 773                if let AgentMessageContent::ToolUse(last_tool_use) = last_content {
 774                    assert_eq!(last_tool_use.name.as_ref(), "word_list");
 775                    if tool_call.status == acp::ToolCallStatus::Pending {
 776                        if !last_tool_use.is_input_complete
 777                            && last_tool_use.input.get("g").is_none()
 778                        {
 779                            saw_partial_tool_use = true;
 780                        }
 781                    } else {
 782                        last_tool_use
 783                            .input
 784                            .get("a")
 785                            .expect("'a' has streamed because input is now complete");
 786                        last_tool_use
 787                            .input
 788                            .get("g")
 789                            .expect("'g' has streamed because input is now complete");
 790                    }
 791                } else {
 792                    panic!("last content should be a tool use");
 793                }
 794            });
 795        }
 796    }
 797
 798    assert!(
 799        saw_partial_tool_use,
 800        "should see at least one partially streamed tool use in the history"
 801    );
 802}
 803
 804#[gpui::test]
 805async fn test_tool_authorization(cx: &mut TestAppContext) {
 806    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 807    let fake_model = model.as_fake();
 808
 809    let mut events = thread
 810        .update(cx, |thread, cx| {
 811            thread.add_tool(ToolRequiringPermission);
 812            thread.send(UserMessageId::new(), ["abc"], cx)
 813        })
 814        .unwrap();
 815    cx.run_until_parked();
 816    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
 817        LanguageModelToolUse {
 818            id: "tool_id_1".into(),
 819            name: ToolRequiringPermission::NAME.into(),
 820            raw_input: "{}".into(),
 821            input: json!({}),
 822            is_input_complete: true,
 823            thought_signature: None,
 824        },
 825    ));
 826    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
 827        LanguageModelToolUse {
 828            id: "tool_id_2".into(),
 829            name: ToolRequiringPermission::NAME.into(),
 830            raw_input: "{}".into(),
 831            input: json!({}),
 832            is_input_complete: true,
 833            thought_signature: None,
 834        },
 835    ));
 836    fake_model.end_last_completion_stream();
 837    let tool_call_auth_1 = next_tool_call_authorization(&mut events).await;
 838    let tool_call_auth_2 = next_tool_call_authorization(&mut events).await;
 839
 840    // Approve the first - send "allow" option_id (UI transforms "once" to "allow")
 841    tool_call_auth_1
 842        .response
 843        .send(acp_thread::SelectedPermissionOutcome::new(
 844            acp::PermissionOptionId::new("allow"),
 845            acp::PermissionOptionKind::AllowOnce,
 846        ))
 847        .unwrap();
 848    cx.run_until_parked();
 849
 850    // Reject the second - send "deny" option_id directly since Deny is now a button
 851    tool_call_auth_2
 852        .response
 853        .send(acp_thread::SelectedPermissionOutcome::new(
 854            acp::PermissionOptionId::new("deny"),
 855            acp::PermissionOptionKind::RejectOnce,
 856        ))
 857        .unwrap();
 858    cx.run_until_parked();
 859
 860    let completion = fake_model.pending_completions().pop().unwrap();
 861    let message = completion.messages.last().unwrap();
 862    assert_eq!(
 863        message.content,
 864        vec![
 865            language_model::MessageContent::ToolResult(LanguageModelToolResult {
 866                tool_use_id: tool_call_auth_1.tool_call.tool_call_id.0.to_string().into(),
 867                tool_name: ToolRequiringPermission::NAME.into(),
 868                is_error: false,
 869                content: "Allowed".into(),
 870                output: Some("Allowed".into())
 871            }),
 872            language_model::MessageContent::ToolResult(LanguageModelToolResult {
 873                tool_use_id: tool_call_auth_2.tool_call.tool_call_id.0.to_string().into(),
 874                tool_name: ToolRequiringPermission::NAME.into(),
 875                is_error: true,
 876                content: "Permission to run tool denied by user".into(),
 877                output: Some("Permission to run tool denied by user".into())
 878            })
 879        ]
 880    );
 881
 882    // Simulate yet another tool call.
 883    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
 884        LanguageModelToolUse {
 885            id: "tool_id_3".into(),
 886            name: ToolRequiringPermission::NAME.into(),
 887            raw_input: "{}".into(),
 888            input: json!({}),
 889            is_input_complete: true,
 890            thought_signature: None,
 891        },
 892    ));
 893    fake_model.end_last_completion_stream();
 894
 895    // Respond by always allowing tools - send transformed option_id
 896    // (UI transforms "always:tool_requiring_permission" to "always_allow:tool_requiring_permission")
 897    let tool_call_auth_3 = next_tool_call_authorization(&mut events).await;
 898    tool_call_auth_3
 899        .response
 900        .send(acp_thread::SelectedPermissionOutcome::new(
 901            acp::PermissionOptionId::new("always_allow:tool_requiring_permission"),
 902            acp::PermissionOptionKind::AllowAlways,
 903        ))
 904        .unwrap();
 905    cx.run_until_parked();
 906    let completion = fake_model.pending_completions().pop().unwrap();
 907    let message = completion.messages.last().unwrap();
 908    assert_eq!(
 909        message.content,
 910        vec![language_model::MessageContent::ToolResult(
 911            LanguageModelToolResult {
 912                tool_use_id: tool_call_auth_3.tool_call.tool_call_id.0.to_string().into(),
 913                tool_name: ToolRequiringPermission::NAME.into(),
 914                is_error: false,
 915                content: "Allowed".into(),
 916                output: Some("Allowed".into())
 917            }
 918        )]
 919    );
 920
 921    // Simulate a final tool call, ensuring we don't trigger authorization.
 922    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
 923        LanguageModelToolUse {
 924            id: "tool_id_4".into(),
 925            name: ToolRequiringPermission::NAME.into(),
 926            raw_input: "{}".into(),
 927            input: json!({}),
 928            is_input_complete: true,
 929            thought_signature: None,
 930        },
 931    ));
 932    fake_model.end_last_completion_stream();
 933    cx.run_until_parked();
 934    let completion = fake_model.pending_completions().pop().unwrap();
 935    let message = completion.messages.last().unwrap();
 936    assert_eq!(
 937        message.content,
 938        vec![language_model::MessageContent::ToolResult(
 939            LanguageModelToolResult {
 940                tool_use_id: "tool_id_4".into(),
 941                tool_name: ToolRequiringPermission::NAME.into(),
 942                is_error: false,
 943                content: "Allowed".into(),
 944                output: Some("Allowed".into())
 945            }
 946        )]
 947    );
 948}
 949
 950#[gpui::test]
 951async fn test_tool_hallucination(cx: &mut TestAppContext) {
 952    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
 953    let fake_model = model.as_fake();
 954
 955    let mut events = thread
 956        .update(cx, |thread, cx| {
 957            thread.send(UserMessageId::new(), ["abc"], cx)
 958        })
 959        .unwrap();
 960    cx.run_until_parked();
 961    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
 962        LanguageModelToolUse {
 963            id: "tool_id_1".into(),
 964            name: "nonexistent_tool".into(),
 965            raw_input: "{}".into(),
 966            input: json!({}),
 967            is_input_complete: true,
 968            thought_signature: None,
 969        },
 970    ));
 971    fake_model.end_last_completion_stream();
 972
 973    let tool_call = expect_tool_call(&mut events).await;
 974    assert_eq!(tool_call.title, "nonexistent_tool");
 975    assert_eq!(tool_call.status, acp::ToolCallStatus::Pending);
 976    let update = expect_tool_call_update_fields(&mut events).await;
 977    assert_eq!(update.fields.status, Some(acp::ToolCallStatus::Failed));
 978}
 979
 980async fn expect_tool_call(events: &mut UnboundedReceiver<Result<ThreadEvent>>) -> acp::ToolCall {
 981    let event = events
 982        .next()
 983        .await
 984        .expect("no tool call authorization event received")
 985        .unwrap();
 986    match event {
 987        ThreadEvent::ToolCall(tool_call) => tool_call,
 988        event => {
 989            panic!("Unexpected event {event:?}");
 990        }
 991    }
 992}
 993
 994async fn expect_tool_call_update_fields(
 995    events: &mut UnboundedReceiver<Result<ThreadEvent>>,
 996) -> acp::ToolCallUpdate {
 997    let event = events
 998        .next()
 999        .await
1000        .expect("no tool call authorization event received")
1001        .unwrap();
1002    match event {
1003        ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(update)) => update,
1004        event => {
1005            panic!("Unexpected event {event:?}");
1006        }
1007    }
1008}
1009
1010async fn expect_plan(events: &mut UnboundedReceiver<Result<ThreadEvent>>) -> acp::Plan {
1011    let event = events
1012        .next()
1013        .await
1014        .expect("no plan event received")
1015        .unwrap();
1016    match event {
1017        ThreadEvent::Plan(plan) => plan,
1018        event => {
1019            panic!("Unexpected event {event:?}");
1020        }
1021    }
1022}
1023
1024async fn next_tool_call_authorization(
1025    events: &mut UnboundedReceiver<Result<ThreadEvent>>,
1026) -> ToolCallAuthorization {
1027    loop {
1028        let event = events
1029            .next()
1030            .await
1031            .expect("no tool call authorization event received")
1032            .unwrap();
1033        if let ThreadEvent::ToolCallAuthorization(tool_call_authorization) = event {
1034            let permission_kinds = tool_call_authorization
1035                .options
1036                .first_option_of_kind(acp::PermissionOptionKind::AllowAlways)
1037                .map(|option| option.kind);
1038            let allow_once = tool_call_authorization
1039                .options
1040                .first_option_of_kind(acp::PermissionOptionKind::AllowOnce)
1041                .map(|option| option.kind);
1042
1043            assert_eq!(
1044                permission_kinds,
1045                Some(acp::PermissionOptionKind::AllowAlways)
1046            );
1047            assert_eq!(allow_once, Some(acp::PermissionOptionKind::AllowOnce));
1048            return tool_call_authorization;
1049        }
1050    }
1051}
1052
1053#[test]
1054fn test_permission_options_terminal_with_pattern() {
1055    let permission_options = ToolPermissionContext::new(
1056        TerminalTool::NAME,
1057        vec!["cargo build --release".to_string()],
1058    )
1059    .build_permission_options();
1060
1061    let PermissionOptions::Dropdown(choices) = permission_options else {
1062        panic!("Expected dropdown permission options");
1063    };
1064
1065    assert_eq!(choices.len(), 3);
1066    let labels: Vec<&str> = choices
1067        .iter()
1068        .map(|choice| choice.allow.name.as_ref())
1069        .collect();
1070    assert!(labels.contains(&"Always for terminal"));
1071    assert!(labels.contains(&"Always for `cargo build` commands"));
1072    assert!(labels.contains(&"Only this time"));
1073}
1074
1075#[test]
1076fn test_permission_options_terminal_command_with_flag_second_token() {
1077    let permission_options =
1078        ToolPermissionContext::new(TerminalTool::NAME, vec!["ls -la".to_string()])
1079            .build_permission_options();
1080
1081    let PermissionOptions::Dropdown(choices) = permission_options else {
1082        panic!("Expected dropdown permission options");
1083    };
1084
1085    assert_eq!(choices.len(), 3);
1086    let labels: Vec<&str> = choices
1087        .iter()
1088        .map(|choice| choice.allow.name.as_ref())
1089        .collect();
1090    assert!(labels.contains(&"Always for terminal"));
1091    assert!(labels.contains(&"Always for `ls` commands"));
1092    assert!(labels.contains(&"Only this time"));
1093}
1094
1095#[test]
1096fn test_permission_options_terminal_single_word_command() {
1097    let permission_options =
1098        ToolPermissionContext::new(TerminalTool::NAME, vec!["whoami".to_string()])
1099            .build_permission_options();
1100
1101    let PermissionOptions::Dropdown(choices) = permission_options else {
1102        panic!("Expected dropdown permission options");
1103    };
1104
1105    assert_eq!(choices.len(), 3);
1106    let labels: Vec<&str> = choices
1107        .iter()
1108        .map(|choice| choice.allow.name.as_ref())
1109        .collect();
1110    assert!(labels.contains(&"Always for terminal"));
1111    assert!(labels.contains(&"Always for `whoami` commands"));
1112    assert!(labels.contains(&"Only this time"));
1113}
1114
1115#[test]
1116fn test_permission_options_edit_file_with_path_pattern() {
1117    let permission_options =
1118        ToolPermissionContext::new(EditFileTool::NAME, vec!["src/main.rs".to_string()])
1119            .build_permission_options();
1120
1121    let PermissionOptions::Dropdown(choices) = permission_options else {
1122        panic!("Expected dropdown permission options");
1123    };
1124
1125    let labels: Vec<&str> = choices
1126        .iter()
1127        .map(|choice| choice.allow.name.as_ref())
1128        .collect();
1129    assert!(labels.contains(&"Always for edit file"));
1130    assert!(labels.contains(&"Always for `src/`"));
1131}
1132
1133#[test]
1134fn test_permission_options_fetch_with_domain_pattern() {
1135    let permission_options =
1136        ToolPermissionContext::new(FetchTool::NAME, vec!["https://docs.rs/gpui".to_string()])
1137            .build_permission_options();
1138
1139    let PermissionOptions::Dropdown(choices) = permission_options else {
1140        panic!("Expected dropdown permission options");
1141    };
1142
1143    let labels: Vec<&str> = choices
1144        .iter()
1145        .map(|choice| choice.allow.name.as_ref())
1146        .collect();
1147    assert!(labels.contains(&"Always for fetch"));
1148    assert!(labels.contains(&"Always for `docs.rs`"));
1149}
1150
1151#[test]
1152fn test_permission_options_without_pattern() {
1153    let permission_options = ToolPermissionContext::new(
1154        TerminalTool::NAME,
1155        vec!["./deploy.sh --production".to_string()],
1156    )
1157    .build_permission_options();
1158
1159    let PermissionOptions::Dropdown(choices) = permission_options else {
1160        panic!("Expected dropdown permission options");
1161    };
1162
1163    assert_eq!(choices.len(), 2);
1164    let labels: Vec<&str> = choices
1165        .iter()
1166        .map(|choice| choice.allow.name.as_ref())
1167        .collect();
1168    assert!(labels.contains(&"Always for terminal"));
1169    assert!(labels.contains(&"Only this time"));
1170    assert!(!labels.iter().any(|label| label.contains("commands")));
1171}
1172
1173#[test]
1174fn test_permission_options_symlink_target_are_flat_once_only() {
1175    let permission_options =
1176        ToolPermissionContext::symlink_target(EditFileTool::NAME, vec!["/outside/file.txt".into()])
1177            .build_permission_options();
1178
1179    let PermissionOptions::Flat(options) = permission_options else {
1180        panic!("Expected flat permission options for symlink target authorization");
1181    };
1182
1183    assert_eq!(options.len(), 2);
1184    assert!(options.iter().any(|option| {
1185        option.option_id.0.as_ref() == "allow"
1186            && option.kind == acp::PermissionOptionKind::AllowOnce
1187    }));
1188    assert!(options.iter().any(|option| {
1189        option.option_id.0.as_ref() == "deny"
1190            && option.kind == acp::PermissionOptionKind::RejectOnce
1191    }));
1192}
1193
1194#[test]
1195fn test_permission_option_ids_for_terminal() {
1196    let permission_options = ToolPermissionContext::new(
1197        TerminalTool::NAME,
1198        vec!["cargo build --release".to_string()],
1199    )
1200    .build_permission_options();
1201
1202    let PermissionOptions::Dropdown(choices) = permission_options else {
1203        panic!("Expected dropdown permission options");
1204    };
1205
1206    // Expect 3 choices: always-tool, always-pattern, once
1207    assert_eq!(choices.len(), 3);
1208
1209    // First two choices both use the tool-level option IDs
1210    assert_eq!(
1211        choices[0].allow.option_id.0.as_ref(),
1212        "always_allow:terminal"
1213    );
1214    assert_eq!(choices[0].deny.option_id.0.as_ref(), "always_deny:terminal");
1215    assert!(choices[0].sub_patterns.is_empty());
1216
1217    assert_eq!(
1218        choices[1].allow.option_id.0.as_ref(),
1219        "always_allow:terminal"
1220    );
1221    assert_eq!(choices[1].deny.option_id.0.as_ref(), "always_deny:terminal");
1222    assert_eq!(choices[1].sub_patterns, vec!["^cargo\\s+build(\\s|$)"]);
1223
1224    // Third choice is the one-time allow/deny
1225    assert_eq!(choices[2].allow.option_id.0.as_ref(), "allow");
1226    assert_eq!(choices[2].deny.option_id.0.as_ref(), "deny");
1227    assert!(choices[2].sub_patterns.is_empty());
1228}
1229
1230#[test]
1231fn test_permission_options_terminal_pipeline_produces_dropdown_with_patterns() {
1232    let permission_options = ToolPermissionContext::new(
1233        TerminalTool::NAME,
1234        vec!["cargo test 2>&1 | tail".to_string()],
1235    )
1236    .build_permission_options();
1237
1238    let PermissionOptions::DropdownWithPatterns {
1239        choices,
1240        patterns,
1241        tool_name,
1242    } = permission_options
1243    else {
1244        panic!("Expected DropdownWithPatterns permission options for pipeline command");
1245    };
1246
1247    assert_eq!(tool_name, TerminalTool::NAME);
1248
1249    // Should have "Always for terminal" and "Only this time" choices
1250    assert_eq!(choices.len(), 2);
1251    let labels: Vec<&str> = choices
1252        .iter()
1253        .map(|choice| choice.allow.name.as_ref())
1254        .collect();
1255    assert!(labels.contains(&"Always for terminal"));
1256    assert!(labels.contains(&"Only this time"));
1257
1258    // Should have per-command patterns for "cargo test" and "tail"
1259    assert_eq!(patterns.len(), 2);
1260    let pattern_names: Vec<&str> = patterns.iter().map(|cp| cp.display_name.as_str()).collect();
1261    assert!(pattern_names.contains(&"cargo test"));
1262    assert!(pattern_names.contains(&"tail"));
1263
1264    // Verify patterns are valid regex patterns
1265    let regex_patterns: Vec<&str> = patterns.iter().map(|cp| cp.pattern.as_str()).collect();
1266    assert!(regex_patterns.contains(&"^cargo\\s+test(\\s|$)"));
1267    assert!(regex_patterns.contains(&"^tail\\b"));
1268}
1269
1270#[test]
1271fn test_permission_options_terminal_pipeline_with_chaining() {
1272    let permission_options = ToolPermissionContext::new(
1273        TerminalTool::NAME,
1274        vec!["npm install && npm test | tail".to_string()],
1275    )
1276    .build_permission_options();
1277
1278    let PermissionOptions::DropdownWithPatterns { patterns, .. } = permission_options else {
1279        panic!("Expected DropdownWithPatterns for chained pipeline command");
1280    };
1281
1282    // With subcommand-aware patterns, "npm install" and "npm test" are distinct
1283    assert_eq!(patterns.len(), 3);
1284    let pattern_names: Vec<&str> = patterns.iter().map(|cp| cp.display_name.as_str()).collect();
1285    assert!(pattern_names.contains(&"npm install"));
1286    assert!(pattern_names.contains(&"npm test"));
1287    assert!(pattern_names.contains(&"tail"));
1288}
1289
1290#[gpui::test]
1291#[cfg_attr(not(feature = "e2e"), ignore)]
1292async fn test_concurrent_tool_calls(cx: &mut TestAppContext) {
1293    let ThreadTest { thread, .. } = setup(cx, TestModel::Sonnet4).await;
1294
1295    // Test concurrent tool calls with different delay times
1296    let events = thread
1297        .update(cx, |thread, cx| {
1298            thread.add_tool(DelayTool);
1299            thread.send(
1300                UserMessageId::new(),
1301                [
1302                    "Call the delay tool twice in the same message.",
1303                    "Once with 100ms. Once with 300ms.",
1304                    "When both timers are complete, describe the outputs.",
1305                ],
1306                cx,
1307            )
1308        })
1309        .unwrap()
1310        .collect()
1311        .await;
1312
1313    let stop_reasons = stop_events(events);
1314    assert_eq!(stop_reasons, vec![acp::StopReason::EndTurn]);
1315
1316    thread.update(cx, |thread, _cx| {
1317        let last_message = thread.last_received_or_pending_message().unwrap();
1318        let agent_message = last_message.as_agent_message().unwrap();
1319        let text = agent_message
1320            .content
1321            .iter()
1322            .filter_map(|content| {
1323                if let AgentMessageContent::Text(text) = content {
1324                    Some(text.as_str())
1325                } else {
1326                    None
1327                }
1328            })
1329            .collect::<String>();
1330
1331        assert!(text.contains("Ding"));
1332    });
1333}
1334
1335#[gpui::test]
1336async fn test_profiles(cx: &mut TestAppContext) {
1337    let ThreadTest {
1338        model, thread, fs, ..
1339    } = setup(cx, TestModel::Fake).await;
1340    let fake_model = model.as_fake();
1341
1342    thread.update(cx, |thread, _cx| {
1343        thread.add_tool(DelayTool);
1344        thread.add_tool(EchoTool);
1345        thread.add_tool(InfiniteTool);
1346    });
1347
1348    // Override profiles and wait for settings to be loaded.
1349    fs.insert_file(
1350        paths::settings_file(),
1351        json!({
1352            "agent": {
1353                "profiles": {
1354                    "test-1": {
1355                        "name": "Test Profile 1",
1356                        "tools": {
1357                            EchoTool::NAME: true,
1358                            DelayTool::NAME: true,
1359                        }
1360                    },
1361                    "test-2": {
1362                        "name": "Test Profile 2",
1363                        "tools": {
1364                            InfiniteTool::NAME: true,
1365                        }
1366                    }
1367                }
1368            }
1369        })
1370        .to_string()
1371        .into_bytes(),
1372    )
1373    .await;
1374    cx.run_until_parked();
1375
1376    // Test that test-1 profile (default) has echo and delay tools
1377    thread
1378        .update(cx, |thread, cx| {
1379            thread.set_profile(AgentProfileId("test-1".into()), cx);
1380            thread.send(UserMessageId::new(), ["test"], cx)
1381        })
1382        .unwrap();
1383    cx.run_until_parked();
1384
1385    let mut pending_completions = fake_model.pending_completions();
1386    assert_eq!(pending_completions.len(), 1);
1387    let completion = pending_completions.pop().unwrap();
1388    let tool_names: Vec<String> = completion
1389        .tools
1390        .iter()
1391        .map(|tool| tool.name.clone())
1392        .collect();
1393    assert_eq!(tool_names, vec![DelayTool::NAME, EchoTool::NAME]);
1394    fake_model.end_last_completion_stream();
1395
1396    // Switch to test-2 profile, and verify that it has only the infinite tool.
1397    thread
1398        .update(cx, |thread, cx| {
1399            thread.set_profile(AgentProfileId("test-2".into()), cx);
1400            thread.send(UserMessageId::new(), ["test2"], cx)
1401        })
1402        .unwrap();
1403    cx.run_until_parked();
1404    let mut pending_completions = fake_model.pending_completions();
1405    assert_eq!(pending_completions.len(), 1);
1406    let completion = pending_completions.pop().unwrap();
1407    let tool_names: Vec<String> = completion
1408        .tools
1409        .iter()
1410        .map(|tool| tool.name.clone())
1411        .collect();
1412    assert_eq!(tool_names, vec![InfiniteTool::NAME]);
1413}
1414
1415#[gpui::test]
1416async fn test_mcp_tools(cx: &mut TestAppContext) {
1417    let ThreadTest {
1418        model,
1419        thread,
1420        context_server_store,
1421        fs,
1422        ..
1423    } = setup(cx, TestModel::Fake).await;
1424    let fake_model = model.as_fake();
1425
1426    // Override profiles and wait for settings to be loaded.
1427    fs.insert_file(
1428        paths::settings_file(),
1429        json!({
1430            "agent": {
1431                "tool_permissions": { "default": "allow" },
1432                "profiles": {
1433                    "test": {
1434                        "name": "Test Profile",
1435                        "enable_all_context_servers": true,
1436                        "tools": {
1437                            EchoTool::NAME: true,
1438                        }
1439                    },
1440                }
1441            }
1442        })
1443        .to_string()
1444        .into_bytes(),
1445    )
1446    .await;
1447    cx.run_until_parked();
1448    thread.update(cx, |thread, cx| {
1449        thread.set_profile(AgentProfileId("test".into()), cx)
1450    });
1451
1452    let mut mcp_tool_calls = setup_context_server(
1453        "test_server",
1454        vec![context_server::types::Tool {
1455            name: "echo".into(),
1456            description: None,
1457            input_schema: serde_json::to_value(EchoTool::input_schema(
1458                LanguageModelToolSchemaFormat::JsonSchema,
1459            ))
1460            .unwrap(),
1461            output_schema: None,
1462            annotations: None,
1463        }],
1464        &context_server_store,
1465        cx,
1466    );
1467
1468    let events = thread.update(cx, |thread, cx| {
1469        thread.send(UserMessageId::new(), ["Hey"], cx).unwrap()
1470    });
1471    cx.run_until_parked();
1472
1473    // Simulate the model calling the MCP tool.
1474    let completion = fake_model.pending_completions().pop().unwrap();
1475    assert_eq!(tool_names_for_completion(&completion), vec!["echo"]);
1476    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
1477        LanguageModelToolUse {
1478            id: "tool_1".into(),
1479            name: "echo".into(),
1480            raw_input: json!({"text": "test"}).to_string(),
1481            input: json!({"text": "test"}),
1482            is_input_complete: true,
1483            thought_signature: None,
1484        },
1485    ));
1486    fake_model.end_last_completion_stream();
1487    cx.run_until_parked();
1488
1489    let (tool_call_params, tool_call_response) = mcp_tool_calls.next().await.unwrap();
1490    assert_eq!(tool_call_params.name, "echo");
1491    assert_eq!(tool_call_params.arguments, Some(json!({"text": "test"})));
1492    tool_call_response
1493        .send(context_server::types::CallToolResponse {
1494            content: vec![context_server::types::ToolResponseContent::Text {
1495                text: "test".into(),
1496            }],
1497            is_error: None,
1498            meta: None,
1499            structured_content: None,
1500        })
1501        .unwrap();
1502    cx.run_until_parked();
1503
1504    assert_eq!(tool_names_for_completion(&completion), vec!["echo"]);
1505    fake_model.send_last_completion_stream_text_chunk("Done!");
1506    fake_model.end_last_completion_stream();
1507    events.collect::<Vec<_>>().await;
1508
1509    // Send again after adding the echo tool, ensuring the name collision is resolved.
1510    let events = thread.update(cx, |thread, cx| {
1511        thread.add_tool(EchoTool);
1512        thread.send(UserMessageId::new(), ["Go"], cx).unwrap()
1513    });
1514    cx.run_until_parked();
1515    let completion = fake_model.pending_completions().pop().unwrap();
1516    assert_eq!(
1517        tool_names_for_completion(&completion),
1518        vec!["echo", "test_server_echo"]
1519    );
1520    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
1521        LanguageModelToolUse {
1522            id: "tool_2".into(),
1523            name: "test_server_echo".into(),
1524            raw_input: json!({"text": "mcp"}).to_string(),
1525            input: json!({"text": "mcp"}),
1526            is_input_complete: true,
1527            thought_signature: None,
1528        },
1529    ));
1530    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
1531        LanguageModelToolUse {
1532            id: "tool_3".into(),
1533            name: "echo".into(),
1534            raw_input: json!({"text": "native"}).to_string(),
1535            input: json!({"text": "native"}),
1536            is_input_complete: true,
1537            thought_signature: None,
1538        },
1539    ));
1540    fake_model.end_last_completion_stream();
1541    cx.run_until_parked();
1542
1543    let (tool_call_params, tool_call_response) = mcp_tool_calls.next().await.unwrap();
1544    assert_eq!(tool_call_params.name, "echo");
1545    assert_eq!(tool_call_params.arguments, Some(json!({"text": "mcp"})));
1546    tool_call_response
1547        .send(context_server::types::CallToolResponse {
1548            content: vec![context_server::types::ToolResponseContent::Text { text: "mcp".into() }],
1549            is_error: None,
1550            meta: None,
1551            structured_content: None,
1552        })
1553        .unwrap();
1554    cx.run_until_parked();
1555
1556    // Ensure the tool results were inserted with the correct names.
1557    let completion = fake_model.pending_completions().pop().unwrap();
1558    assert_eq!(
1559        completion.messages.last().unwrap().content,
1560        vec![
1561            MessageContent::ToolResult(LanguageModelToolResult {
1562                tool_use_id: "tool_3".into(),
1563                tool_name: "echo".into(),
1564                is_error: false,
1565                content: "native".into(),
1566                output: Some("native".into()),
1567            },),
1568            MessageContent::ToolResult(LanguageModelToolResult {
1569                tool_use_id: "tool_2".into(),
1570                tool_name: "test_server_echo".into(),
1571                is_error: false,
1572                content: "mcp".into(),
1573                output: Some("mcp".into()),
1574            },),
1575        ]
1576    );
1577    fake_model.end_last_completion_stream();
1578    events.collect::<Vec<_>>().await;
1579}
1580
1581#[gpui::test]
1582async fn test_mcp_tool_result_displayed_when_server_disconnected(cx: &mut TestAppContext) {
1583    let ThreadTest {
1584        model,
1585        thread,
1586        context_server_store,
1587        fs,
1588        ..
1589    } = setup(cx, TestModel::Fake).await;
1590    let fake_model = model.as_fake();
1591
1592    // Setup settings to allow MCP tools
1593    fs.insert_file(
1594        paths::settings_file(),
1595        json!({
1596            "agent": {
1597                "always_allow_tool_actions": true,
1598                "profiles": {
1599                    "test": {
1600                        "name": "Test Profile",
1601                        "enable_all_context_servers": true,
1602                        "tools": {}
1603                    },
1604                }
1605            }
1606        })
1607        .to_string()
1608        .into_bytes(),
1609    )
1610    .await;
1611    cx.run_until_parked();
1612    thread.update(cx, |thread, cx| {
1613        thread.set_profile(AgentProfileId("test".into()), cx)
1614    });
1615
1616    // Setup a context server with a tool
1617    let mut mcp_tool_calls = setup_context_server(
1618        "github_server",
1619        vec![context_server::types::Tool {
1620            name: "issue_read".into(),
1621            description: Some("Read a GitHub issue".into()),
1622            input_schema: json!({
1623                "type": "object",
1624                "properties": {
1625                    "issue_url": { "type": "string" }
1626                }
1627            }),
1628            output_schema: None,
1629            annotations: None,
1630        }],
1631        &context_server_store,
1632        cx,
1633    );
1634
1635    // Send a message and have the model call the MCP tool
1636    let events = thread.update(cx, |thread, cx| {
1637        thread
1638            .send(UserMessageId::new(), ["Read issue #47404"], cx)
1639            .unwrap()
1640    });
1641    cx.run_until_parked();
1642
1643    // Verify the MCP tool is available to the model
1644    let completion = fake_model.pending_completions().pop().unwrap();
1645    assert_eq!(
1646        tool_names_for_completion(&completion),
1647        vec!["issue_read"],
1648        "MCP tool should be available"
1649    );
1650
1651    // Simulate the model calling the MCP tool
1652    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
1653        LanguageModelToolUse {
1654            id: "tool_1".into(),
1655            name: "issue_read".into(),
1656            raw_input: json!({"issue_url": "https://github.com/zed-industries/zed/issues/47404"})
1657                .to_string(),
1658            input: json!({"issue_url": "https://github.com/zed-industries/zed/issues/47404"}),
1659            is_input_complete: true,
1660            thought_signature: None,
1661        },
1662    ));
1663    fake_model.end_last_completion_stream();
1664    cx.run_until_parked();
1665
1666    // The MCP server receives the tool call and responds with content
1667    let expected_tool_output = "Issue #47404: Tool call results are cleared upon app close";
1668    let (tool_call_params, tool_call_response) = mcp_tool_calls.next().await.unwrap();
1669    assert_eq!(tool_call_params.name, "issue_read");
1670    tool_call_response
1671        .send(context_server::types::CallToolResponse {
1672            content: vec![context_server::types::ToolResponseContent::Text {
1673                text: expected_tool_output.into(),
1674            }],
1675            is_error: None,
1676            meta: None,
1677            structured_content: None,
1678        })
1679        .unwrap();
1680    cx.run_until_parked();
1681
1682    // After tool completes, the model continues with a new completion request
1683    // that includes the tool results. We need to respond to this.
1684    let _completion = fake_model.pending_completions().pop().unwrap();
1685    fake_model.send_last_completion_stream_text_chunk("I found the issue!");
1686    fake_model
1687        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
1688    fake_model.end_last_completion_stream();
1689    events.collect::<Vec<_>>().await;
1690
1691    // Verify the tool result is stored in the thread by checking the markdown output.
1692    // The tool result is in the first assistant message (not the last one, which is
1693    // the model's response after the tool completed).
1694    thread.update(cx, |thread, _cx| {
1695        let markdown = thread.to_markdown();
1696        assert!(
1697            markdown.contains("**Tool Result**: issue_read"),
1698            "Thread should contain tool result header"
1699        );
1700        assert!(
1701            markdown.contains(expected_tool_output),
1702            "Thread should contain tool output: {}",
1703            expected_tool_output
1704        );
1705    });
1706
1707    // Simulate app restart: disconnect the MCP server.
1708    // After restart, the MCP server won't be connected yet when the thread is replayed.
1709    context_server_store.update(cx, |store, cx| {
1710        let _ = store.stop_server(&ContextServerId("github_server".into()), cx);
1711    });
1712    cx.run_until_parked();
1713
1714    // Replay the thread (this is what happens when loading a saved thread)
1715    let mut replay_events = thread.update(cx, |thread, cx| thread.replay(cx));
1716
1717    let mut found_tool_call = None;
1718    let mut found_tool_call_update_with_output = None;
1719
1720    while let Some(event) = replay_events.next().await {
1721        let event = event.unwrap();
1722        match &event {
1723            ThreadEvent::ToolCall(tc) if tc.tool_call_id.to_string() == "tool_1" => {
1724                found_tool_call = Some(tc.clone());
1725            }
1726            ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(update))
1727                if update.tool_call_id.to_string() == "tool_1" =>
1728            {
1729                if update.fields.raw_output.is_some() {
1730                    found_tool_call_update_with_output = Some(update.clone());
1731                }
1732            }
1733            _ => {}
1734        }
1735    }
1736
1737    // The tool call should be found
1738    assert!(
1739        found_tool_call.is_some(),
1740        "Tool call should be emitted during replay"
1741    );
1742
1743    assert!(
1744        found_tool_call_update_with_output.is_some(),
1745        "ToolCallUpdate with raw_output should be emitted even when MCP server is disconnected."
1746    );
1747
1748    let update = found_tool_call_update_with_output.unwrap();
1749    assert_eq!(
1750        update.fields.raw_output,
1751        Some(expected_tool_output.into()),
1752        "raw_output should contain the saved tool result"
1753    );
1754
1755    // Also verify the status is correct (completed, not failed)
1756    assert_eq!(
1757        update.fields.status,
1758        Some(acp::ToolCallStatus::Completed),
1759        "Tool call status should reflect the original completion status"
1760    );
1761}
1762
1763#[gpui::test]
1764async fn test_mcp_tool_truncation(cx: &mut TestAppContext) {
1765    let ThreadTest {
1766        model,
1767        thread,
1768        context_server_store,
1769        fs,
1770        ..
1771    } = setup(cx, TestModel::Fake).await;
1772    let fake_model = model.as_fake();
1773
1774    // Set up a profile with all tools enabled
1775    fs.insert_file(
1776        paths::settings_file(),
1777        json!({
1778            "agent": {
1779                "profiles": {
1780                    "test": {
1781                        "name": "Test Profile",
1782                        "enable_all_context_servers": true,
1783                        "tools": {
1784                            EchoTool::NAME: true,
1785                            DelayTool::NAME: true,
1786                            WordListTool::NAME: true,
1787                            ToolRequiringPermission::NAME: true,
1788                            InfiniteTool::NAME: true,
1789                        }
1790                    },
1791                }
1792            }
1793        })
1794        .to_string()
1795        .into_bytes(),
1796    )
1797    .await;
1798    cx.run_until_parked();
1799
1800    thread.update(cx, |thread, cx| {
1801        thread.set_profile(AgentProfileId("test".into()), cx);
1802        thread.add_tool(EchoTool);
1803        thread.add_tool(DelayTool);
1804        thread.add_tool(WordListTool);
1805        thread.add_tool(ToolRequiringPermission);
1806        thread.add_tool(InfiniteTool);
1807    });
1808
1809    // Set up multiple context servers with some overlapping tool names
1810    let _server1_calls = setup_context_server(
1811        "xxx",
1812        vec![
1813            context_server::types::Tool {
1814                name: "echo".into(), // Conflicts with native EchoTool
1815                description: None,
1816                input_schema: serde_json::to_value(EchoTool::input_schema(
1817                    LanguageModelToolSchemaFormat::JsonSchema,
1818                ))
1819                .unwrap(),
1820                output_schema: None,
1821                annotations: None,
1822            },
1823            context_server::types::Tool {
1824                name: "unique_tool_1".into(),
1825                description: None,
1826                input_schema: json!({"type": "object", "properties": {}}),
1827                output_schema: None,
1828                annotations: None,
1829            },
1830        ],
1831        &context_server_store,
1832        cx,
1833    );
1834
1835    let _server2_calls = setup_context_server(
1836        "yyy",
1837        vec![
1838            context_server::types::Tool {
1839                name: "echo".into(), // Also conflicts with native EchoTool
1840                description: None,
1841                input_schema: serde_json::to_value(EchoTool::input_schema(
1842                    LanguageModelToolSchemaFormat::JsonSchema,
1843                ))
1844                .unwrap(),
1845                output_schema: None,
1846                annotations: None,
1847            },
1848            context_server::types::Tool {
1849                name: "unique_tool_2".into(),
1850                description: None,
1851                input_schema: json!({"type": "object", "properties": {}}),
1852                output_schema: None,
1853                annotations: None,
1854            },
1855            context_server::types::Tool {
1856                name: "a".repeat(MAX_TOOL_NAME_LENGTH - 2),
1857                description: None,
1858                input_schema: json!({"type": "object", "properties": {}}),
1859                output_schema: None,
1860                annotations: None,
1861            },
1862            context_server::types::Tool {
1863                name: "b".repeat(MAX_TOOL_NAME_LENGTH - 1),
1864                description: None,
1865                input_schema: json!({"type": "object", "properties": {}}),
1866                output_schema: None,
1867                annotations: None,
1868            },
1869        ],
1870        &context_server_store,
1871        cx,
1872    );
1873    let _server3_calls = setup_context_server(
1874        "zzz",
1875        vec![
1876            context_server::types::Tool {
1877                name: "a".repeat(MAX_TOOL_NAME_LENGTH - 2),
1878                description: None,
1879                input_schema: json!({"type": "object", "properties": {}}),
1880                output_schema: None,
1881                annotations: None,
1882            },
1883            context_server::types::Tool {
1884                name: "b".repeat(MAX_TOOL_NAME_LENGTH - 1),
1885                description: None,
1886                input_schema: json!({"type": "object", "properties": {}}),
1887                output_schema: None,
1888                annotations: None,
1889            },
1890            context_server::types::Tool {
1891                name: "c".repeat(MAX_TOOL_NAME_LENGTH + 1),
1892                description: None,
1893                input_schema: json!({"type": "object", "properties": {}}),
1894                output_schema: None,
1895                annotations: None,
1896            },
1897        ],
1898        &context_server_store,
1899        cx,
1900    );
1901
1902    // Server with spaces in name - tests snake_case conversion for API compatibility
1903    let _server4_calls = setup_context_server(
1904        "Azure DevOps",
1905        vec![context_server::types::Tool {
1906            name: "echo".into(), // Also conflicts - will be disambiguated as azure_dev_ops_echo
1907            description: None,
1908            input_schema: serde_json::to_value(EchoTool::input_schema(
1909                LanguageModelToolSchemaFormat::JsonSchema,
1910            ))
1911            .unwrap(),
1912            output_schema: None,
1913            annotations: None,
1914        }],
1915        &context_server_store,
1916        cx,
1917    );
1918
1919    thread
1920        .update(cx, |thread, cx| {
1921            thread.send(UserMessageId::new(), ["Go"], cx)
1922        })
1923        .unwrap();
1924    cx.run_until_parked();
1925    let completion = fake_model.pending_completions().pop().unwrap();
1926    assert_eq!(
1927        tool_names_for_completion(&completion),
1928        vec![
1929            "azure_dev_ops_echo",
1930            "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
1931            "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc",
1932            "delay",
1933            "echo",
1934            "infinite",
1935            "tool_requiring_permission",
1936            "unique_tool_1",
1937            "unique_tool_2",
1938            "word_list",
1939            "xxx_echo",
1940            "y_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
1941            "yyy_echo",
1942            "z_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
1943        ]
1944    );
1945}
1946
1947#[gpui::test]
1948#[cfg_attr(not(feature = "e2e"), ignore)]
1949async fn test_cancellation(cx: &mut TestAppContext) {
1950    let ThreadTest { thread, .. } = setup(cx, TestModel::Sonnet4).await;
1951
1952    let mut events = thread
1953        .update(cx, |thread, cx| {
1954            thread.add_tool(InfiniteTool);
1955            thread.add_tool(EchoTool);
1956            thread.send(
1957                UserMessageId::new(),
1958                ["Call the echo tool, then call the infinite tool, then explain their output"],
1959                cx,
1960            )
1961        })
1962        .unwrap();
1963
1964    // Wait until both tools are called.
1965    let mut expected_tools = vec!["Echo", "Infinite Tool"];
1966    let mut echo_id = None;
1967    let mut echo_completed = false;
1968    while let Some(event) = events.next().await {
1969        match event.unwrap() {
1970            ThreadEvent::ToolCall(tool_call) => {
1971                assert_eq!(tool_call.title, expected_tools.remove(0));
1972                if tool_call.title == "Echo" {
1973                    echo_id = Some(tool_call.tool_call_id);
1974                }
1975            }
1976            ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(
1977                acp::ToolCallUpdate {
1978                    tool_call_id,
1979                    fields:
1980                        acp::ToolCallUpdateFields {
1981                            status: Some(acp::ToolCallStatus::Completed),
1982                            ..
1983                        },
1984                    ..
1985                },
1986            )) if Some(&tool_call_id) == echo_id.as_ref() => {
1987                echo_completed = true;
1988            }
1989            _ => {}
1990        }
1991
1992        if expected_tools.is_empty() && echo_completed {
1993            break;
1994        }
1995    }
1996
1997    // Cancel the current send and ensure that the event stream is closed, even
1998    // if one of the tools is still running.
1999    thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2000    let events = events.collect::<Vec<_>>().await;
2001    let last_event = events.last();
2002    assert!(
2003        matches!(
2004            last_event,
2005            Some(Ok(ThreadEvent::Stop(acp::StopReason::Cancelled)))
2006        ),
2007        "unexpected event {last_event:?}"
2008    );
2009
2010    // Ensure we can still send a new message after cancellation.
2011    let events = thread
2012        .update(cx, |thread, cx| {
2013            thread.send(
2014                UserMessageId::new(),
2015                ["Testing: reply with 'Hello' then stop."],
2016                cx,
2017            )
2018        })
2019        .unwrap()
2020        .collect::<Vec<_>>()
2021        .await;
2022    thread.update(cx, |thread, _cx| {
2023        let message = thread.last_received_or_pending_message().unwrap();
2024        let agent_message = message.as_agent_message().unwrap();
2025        assert_eq!(
2026            agent_message.content,
2027            vec![AgentMessageContent::Text("Hello".to_string())]
2028        );
2029    });
2030    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
2031}
2032
2033#[gpui::test]
2034async fn test_terminal_tool_cancellation_captures_output(cx: &mut TestAppContext) {
2035    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2036    always_allow_tools(cx);
2037    let fake_model = model.as_fake();
2038
2039    let environment = Rc::new(cx.update(|cx| {
2040        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
2041    }));
2042    let handle = environment.terminal_handle.clone().unwrap();
2043
2044    let mut events = thread
2045        .update(cx, |thread, cx| {
2046            thread.add_tool(crate::TerminalTool::new(
2047                thread.project().clone(),
2048                environment,
2049            ));
2050            thread.send(UserMessageId::new(), ["run a command"], cx)
2051        })
2052        .unwrap();
2053
2054    cx.run_until_parked();
2055
2056    // Simulate the model calling the terminal tool
2057    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2058        LanguageModelToolUse {
2059            id: "terminal_tool_1".into(),
2060            name: TerminalTool::NAME.into(),
2061            raw_input: r#"{"command": "sleep 1000", "cd": "."}"#.into(),
2062            input: json!({"command": "sleep 1000", "cd": "."}),
2063            is_input_complete: true,
2064            thought_signature: None,
2065        },
2066    ));
2067    fake_model.end_last_completion_stream();
2068
2069    // Wait for the terminal tool to start running
2070    wait_for_terminal_tool_started(&mut events, cx).await;
2071
2072    // Cancel the thread while the terminal is running
2073    thread.update(cx, |thread, cx| thread.cancel(cx)).detach();
2074
2075    // Collect remaining events, driving the executor to let cancellation complete
2076    let remaining_events = collect_events_until_stop(&mut events, cx).await;
2077
2078    // Verify the terminal was killed
2079    assert!(
2080        handle.was_killed(),
2081        "expected terminal handle to be killed on cancellation"
2082    );
2083
2084    // Verify we got a cancellation stop event
2085    assert_eq!(
2086        stop_events(remaining_events),
2087        vec![acp::StopReason::Cancelled],
2088    );
2089
2090    // Verify the tool result contains the terminal output, not just "Tool canceled by user"
2091    thread.update(cx, |thread, _cx| {
2092        let message = thread.last_received_or_pending_message().unwrap();
2093        let agent_message = message.as_agent_message().unwrap();
2094
2095        let tool_use = agent_message
2096            .content
2097            .iter()
2098            .find_map(|content| match content {
2099                AgentMessageContent::ToolUse(tool_use) => Some(tool_use),
2100                _ => None,
2101            })
2102            .expect("expected tool use in agent message");
2103
2104        let tool_result = agent_message
2105            .tool_results
2106            .get(&tool_use.id)
2107            .expect("expected tool result");
2108
2109        let result_text = match &tool_result.content {
2110            language_model::LanguageModelToolResultContent::Text(text) => text.to_string(),
2111            _ => panic!("expected text content in tool result"),
2112        };
2113
2114        // "partial output" comes from FakeTerminalHandle's output field
2115        assert!(
2116            result_text.contains("partial output"),
2117            "expected tool result to contain terminal output, got: {result_text}"
2118        );
2119        // Match the actual format from process_content in terminal_tool.rs
2120        assert!(
2121            result_text.contains("The user stopped this command"),
2122            "expected tool result to indicate user stopped, got: {result_text}"
2123        );
2124    });
2125
2126    // Verify we can send a new message after cancellation
2127    verify_thread_recovery(&thread, &fake_model, cx).await;
2128}
2129
2130#[gpui::test]
2131async fn test_cancellation_aware_tool_responds_to_cancellation(cx: &mut TestAppContext) {
2132    // This test verifies that tools which properly handle cancellation via
2133    // `event_stream.cancelled_by_user()` (like edit_file_tool) respond promptly
2134    // to cancellation and report that they were cancelled.
2135    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2136    always_allow_tools(cx);
2137    let fake_model = model.as_fake();
2138
2139    let (tool, was_cancelled) = CancellationAwareTool::new();
2140
2141    let mut events = thread
2142        .update(cx, |thread, cx| {
2143            thread.add_tool(tool);
2144            thread.send(
2145                UserMessageId::new(),
2146                ["call the cancellation aware tool"],
2147                cx,
2148            )
2149        })
2150        .unwrap();
2151
2152    cx.run_until_parked();
2153
2154    // Simulate the model calling the cancellation-aware tool
2155    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2156        LanguageModelToolUse {
2157            id: "cancellation_aware_1".into(),
2158            name: "cancellation_aware".into(),
2159            raw_input: r#"{}"#.into(),
2160            input: json!({}),
2161            is_input_complete: true,
2162            thought_signature: None,
2163        },
2164    ));
2165    fake_model.end_last_completion_stream();
2166
2167    cx.run_until_parked();
2168
2169    // Wait for the tool call to be reported
2170    let mut tool_started = false;
2171    let deadline = cx.executor().num_cpus() * 100;
2172    for _ in 0..deadline {
2173        cx.run_until_parked();
2174
2175        while let Some(Some(event)) = events.next().now_or_never() {
2176            if let Ok(ThreadEvent::ToolCall(tool_call)) = &event {
2177                if tool_call.title == "Cancellation Aware Tool" {
2178                    tool_started = true;
2179                    break;
2180                }
2181            }
2182        }
2183
2184        if tool_started {
2185            break;
2186        }
2187
2188        cx.background_executor
2189            .timer(Duration::from_millis(10))
2190            .await;
2191    }
2192    assert!(tool_started, "expected cancellation aware tool to start");
2193
2194    // Cancel the thread and wait for it to complete
2195    let cancel_task = thread.update(cx, |thread, cx| thread.cancel(cx));
2196
2197    // The cancel task should complete promptly because the tool handles cancellation
2198    let timeout = cx.background_executor.timer(Duration::from_secs(5));
2199    futures::select! {
2200        _ = cancel_task.fuse() => {}
2201        _ = timeout.fuse() => {
2202            panic!("cancel task timed out - tool did not respond to cancellation");
2203        }
2204    }
2205
2206    // Verify the tool detected cancellation via its flag
2207    assert!(
2208        was_cancelled.load(std::sync::atomic::Ordering::SeqCst),
2209        "tool should have detected cancellation via event_stream.cancelled_by_user()"
2210    );
2211
2212    // Collect remaining events
2213    let remaining_events = collect_events_until_stop(&mut events, cx).await;
2214
2215    // Verify we got a cancellation stop event
2216    assert_eq!(
2217        stop_events(remaining_events),
2218        vec![acp::StopReason::Cancelled],
2219    );
2220
2221    // Verify we can send a new message after cancellation
2222    verify_thread_recovery(&thread, &fake_model, cx).await;
2223}
2224
2225/// Helper to verify thread can recover after cancellation by sending a simple message.
2226async fn verify_thread_recovery(
2227    thread: &Entity<Thread>,
2228    fake_model: &FakeLanguageModel,
2229    cx: &mut TestAppContext,
2230) {
2231    let events = thread
2232        .update(cx, |thread, cx| {
2233            thread.send(
2234                UserMessageId::new(),
2235                ["Testing: reply with 'Hello' then stop."],
2236                cx,
2237            )
2238        })
2239        .unwrap();
2240    cx.run_until_parked();
2241    fake_model.send_last_completion_stream_text_chunk("Hello");
2242    fake_model
2243        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2244    fake_model.end_last_completion_stream();
2245
2246    let events = events.collect::<Vec<_>>().await;
2247    thread.update(cx, |thread, _cx| {
2248        let message = thread.last_received_or_pending_message().unwrap();
2249        let agent_message = message.as_agent_message().unwrap();
2250        assert_eq!(
2251            agent_message.content,
2252            vec![AgentMessageContent::Text("Hello".to_string())]
2253        );
2254    });
2255    assert_eq!(stop_events(events), vec![acp::StopReason::EndTurn]);
2256}
2257
2258/// Waits for a terminal tool to start by watching for a ToolCallUpdate with terminal content.
2259async fn wait_for_terminal_tool_started(
2260    events: &mut mpsc::UnboundedReceiver<Result<ThreadEvent>>,
2261    cx: &mut TestAppContext,
2262) {
2263    let deadline = cx.executor().num_cpus() * 100; // Scale with available parallelism
2264    for _ in 0..deadline {
2265        cx.run_until_parked();
2266
2267        while let Some(Some(event)) = events.next().now_or_never() {
2268            if let Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(
2269                update,
2270            ))) = &event
2271            {
2272                if update.fields.content.as_ref().is_some_and(|content| {
2273                    content
2274                        .iter()
2275                        .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))
2276                }) {
2277                    return;
2278                }
2279            }
2280        }
2281
2282        cx.background_executor
2283            .timer(Duration::from_millis(10))
2284            .await;
2285    }
2286    panic!("terminal tool did not start within the expected time");
2287}
2288
2289/// Collects events until a Stop event is received, driving the executor to completion.
2290async fn collect_events_until_stop(
2291    events: &mut mpsc::UnboundedReceiver<Result<ThreadEvent>>,
2292    cx: &mut TestAppContext,
2293) -> Vec<Result<ThreadEvent>> {
2294    let mut collected = Vec::new();
2295    let deadline = cx.executor().num_cpus() * 200;
2296
2297    for _ in 0..deadline {
2298        cx.executor().advance_clock(Duration::from_millis(10));
2299        cx.run_until_parked();
2300
2301        while let Some(Some(event)) = events.next().now_or_never() {
2302            let is_stop = matches!(&event, Ok(ThreadEvent::Stop(_)));
2303            collected.push(event);
2304            if is_stop {
2305                return collected;
2306            }
2307        }
2308    }
2309    panic!(
2310        "did not receive Stop event within the expected time; collected {} events",
2311        collected.len()
2312    );
2313}
2314
2315#[gpui::test]
2316async fn test_truncate_while_terminal_tool_running(cx: &mut TestAppContext) {
2317    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2318    always_allow_tools(cx);
2319    let fake_model = model.as_fake();
2320
2321    let environment = Rc::new(cx.update(|cx| {
2322        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
2323    }));
2324    let handle = environment.terminal_handle.clone().unwrap();
2325
2326    let message_id = UserMessageId::new();
2327    let mut events = thread
2328        .update(cx, |thread, cx| {
2329            thread.add_tool(crate::TerminalTool::new(
2330                thread.project().clone(),
2331                environment,
2332            ));
2333            thread.send(message_id.clone(), ["run a command"], cx)
2334        })
2335        .unwrap();
2336
2337    cx.run_until_parked();
2338
2339    // Simulate the model calling the terminal tool
2340    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2341        LanguageModelToolUse {
2342            id: "terminal_tool_1".into(),
2343            name: TerminalTool::NAME.into(),
2344            raw_input: r#"{"command": "sleep 1000", "cd": "."}"#.into(),
2345            input: json!({"command": "sleep 1000", "cd": "."}),
2346            is_input_complete: true,
2347            thought_signature: None,
2348        },
2349    ));
2350    fake_model.end_last_completion_stream();
2351
2352    // Wait for the terminal tool to start running
2353    wait_for_terminal_tool_started(&mut events, cx).await;
2354
2355    // Truncate the thread while the terminal is running
2356    thread
2357        .update(cx, |thread, cx| thread.truncate(message_id, cx))
2358        .unwrap();
2359
2360    // Drive the executor to let cancellation complete
2361    let _ = collect_events_until_stop(&mut events, cx).await;
2362
2363    // Verify the terminal was killed
2364    assert!(
2365        handle.was_killed(),
2366        "expected terminal handle to be killed on truncate"
2367    );
2368
2369    // Verify the thread is empty after truncation
2370    thread.update(cx, |thread, _cx| {
2371        assert_eq!(
2372            thread.to_markdown(),
2373            "",
2374            "expected thread to be empty after truncating the only message"
2375        );
2376    });
2377
2378    // Verify we can send a new message after truncation
2379    verify_thread_recovery(&thread, &fake_model, cx).await;
2380}
2381
2382#[gpui::test]
2383async fn test_cancel_multiple_concurrent_terminal_tools(cx: &mut TestAppContext) {
2384    // Tests that cancellation properly kills all running terminal tools when multiple are active.
2385    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2386    always_allow_tools(cx);
2387    let fake_model = model.as_fake();
2388
2389    let environment = Rc::new(MultiTerminalEnvironment::new());
2390
2391    let mut events = thread
2392        .update(cx, |thread, cx| {
2393            thread.add_tool(crate::TerminalTool::new(
2394                thread.project().clone(),
2395                environment.clone(),
2396            ));
2397            thread.send(UserMessageId::new(), ["run multiple commands"], cx)
2398        })
2399        .unwrap();
2400
2401    cx.run_until_parked();
2402
2403    // Simulate the model calling two terminal tools
2404    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2405        LanguageModelToolUse {
2406            id: "terminal_tool_1".into(),
2407            name: TerminalTool::NAME.into(),
2408            raw_input: r#"{"command": "sleep 1000", "cd": "."}"#.into(),
2409            input: json!({"command": "sleep 1000", "cd": "."}),
2410            is_input_complete: true,
2411            thought_signature: None,
2412        },
2413    ));
2414    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2415        LanguageModelToolUse {
2416            id: "terminal_tool_2".into(),
2417            name: TerminalTool::NAME.into(),
2418            raw_input: r#"{"command": "sleep 2000", "cd": "."}"#.into(),
2419            input: json!({"command": "sleep 2000", "cd": "."}),
2420            is_input_complete: true,
2421            thought_signature: None,
2422        },
2423    ));
2424    fake_model.end_last_completion_stream();
2425
2426    // Wait for both terminal tools to start by counting terminal content updates
2427    let mut terminals_started = 0;
2428    let deadline = cx.executor().num_cpus() * 100;
2429    for _ in 0..deadline {
2430        cx.run_until_parked();
2431
2432        while let Some(Some(event)) = events.next().now_or_never() {
2433            if let Ok(ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(
2434                update,
2435            ))) = &event
2436            {
2437                if update.fields.content.as_ref().is_some_and(|content| {
2438                    content
2439                        .iter()
2440                        .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))
2441                }) {
2442                    terminals_started += 1;
2443                    if terminals_started >= 2 {
2444                        break;
2445                    }
2446                }
2447            }
2448        }
2449        if terminals_started >= 2 {
2450            break;
2451        }
2452
2453        cx.background_executor
2454            .timer(Duration::from_millis(10))
2455            .await;
2456    }
2457    assert!(
2458        terminals_started >= 2,
2459        "expected 2 terminal tools to start, got {terminals_started}"
2460    );
2461
2462    // Cancel the thread while both terminals are running
2463    thread.update(cx, |thread, cx| thread.cancel(cx)).detach();
2464
2465    // Collect remaining events
2466    let remaining_events = collect_events_until_stop(&mut events, cx).await;
2467
2468    // Verify both terminal handles were killed
2469    let handles = environment.handles();
2470    assert_eq!(
2471        handles.len(),
2472        2,
2473        "expected 2 terminal handles to be created"
2474    );
2475    assert!(
2476        handles[0].was_killed(),
2477        "expected first terminal handle to be killed on cancellation"
2478    );
2479    assert!(
2480        handles[1].was_killed(),
2481        "expected second terminal handle to be killed on cancellation"
2482    );
2483
2484    // Verify we got a cancellation stop event
2485    assert_eq!(
2486        stop_events(remaining_events),
2487        vec![acp::StopReason::Cancelled],
2488    );
2489}
2490
2491#[gpui::test]
2492async fn test_terminal_tool_stopped_via_terminal_card_button(cx: &mut TestAppContext) {
2493    // Tests that clicking the stop button on the terminal card (as opposed to the main
2494    // cancel button) properly reports user stopped via the was_stopped_by_user path.
2495    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2496    always_allow_tools(cx);
2497    let fake_model = model.as_fake();
2498
2499    let environment = Rc::new(cx.update(|cx| {
2500        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
2501    }));
2502    let handle = environment.terminal_handle.clone().unwrap();
2503
2504    let mut events = thread
2505        .update(cx, |thread, cx| {
2506            thread.add_tool(crate::TerminalTool::new(
2507                thread.project().clone(),
2508                environment,
2509            ));
2510            thread.send(UserMessageId::new(), ["run a command"], cx)
2511        })
2512        .unwrap();
2513
2514    cx.run_until_parked();
2515
2516    // Simulate the model calling the terminal tool
2517    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2518        LanguageModelToolUse {
2519            id: "terminal_tool_1".into(),
2520            name: TerminalTool::NAME.into(),
2521            raw_input: r#"{"command": "sleep 1000", "cd": "."}"#.into(),
2522            input: json!({"command": "sleep 1000", "cd": "."}),
2523            is_input_complete: true,
2524            thought_signature: None,
2525        },
2526    ));
2527    fake_model.end_last_completion_stream();
2528
2529    // Wait for the terminal tool to start running
2530    wait_for_terminal_tool_started(&mut events, cx).await;
2531
2532    // Simulate user clicking stop on the terminal card itself.
2533    // This sets the flag and signals exit (simulating what the real UI would do).
2534    handle.set_stopped_by_user(true);
2535    handle.killed.store(true, Ordering::SeqCst);
2536    handle.signal_exit();
2537
2538    // Wait for the tool to complete
2539    cx.run_until_parked();
2540
2541    // The thread continues after tool completion - simulate the model ending its turn
2542    fake_model
2543        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2544    fake_model.end_last_completion_stream();
2545
2546    // Collect remaining events
2547    let remaining_events = collect_events_until_stop(&mut events, cx).await;
2548
2549    // Verify we got an EndTurn (not Cancelled, since we didn't cancel the thread)
2550    assert_eq!(
2551        stop_events(remaining_events),
2552        vec![acp::StopReason::EndTurn],
2553    );
2554
2555    // Verify the tool result indicates user stopped
2556    thread.update(cx, |thread, _cx| {
2557        let message = thread.last_received_or_pending_message().unwrap();
2558        let agent_message = message.as_agent_message().unwrap();
2559
2560        let tool_use = agent_message
2561            .content
2562            .iter()
2563            .find_map(|content| match content {
2564                AgentMessageContent::ToolUse(tool_use) => Some(tool_use),
2565                _ => None,
2566            })
2567            .expect("expected tool use in agent message");
2568
2569        let tool_result = agent_message
2570            .tool_results
2571            .get(&tool_use.id)
2572            .expect("expected tool result");
2573
2574        let result_text = match &tool_result.content {
2575            language_model::LanguageModelToolResultContent::Text(text) => text.to_string(),
2576            _ => panic!("expected text content in tool result"),
2577        };
2578
2579        assert!(
2580            result_text.contains("The user stopped this command"),
2581            "expected tool result to indicate user stopped, got: {result_text}"
2582        );
2583    });
2584}
2585
2586#[gpui::test]
2587async fn test_terminal_tool_timeout_expires(cx: &mut TestAppContext) {
2588    // Tests that when a timeout is configured and expires, the tool result indicates timeout.
2589    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2590    always_allow_tools(cx);
2591    let fake_model = model.as_fake();
2592
2593    let environment = Rc::new(cx.update(|cx| {
2594        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
2595    }));
2596    let handle = environment.terminal_handle.clone().unwrap();
2597
2598    let mut events = thread
2599        .update(cx, |thread, cx| {
2600            thread.add_tool(crate::TerminalTool::new(
2601                thread.project().clone(),
2602                environment,
2603            ));
2604            thread.send(UserMessageId::new(), ["run a command with timeout"], cx)
2605        })
2606        .unwrap();
2607
2608    cx.run_until_parked();
2609
2610    // Simulate the model calling the terminal tool with a short timeout
2611    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
2612        LanguageModelToolUse {
2613            id: "terminal_tool_1".into(),
2614            name: TerminalTool::NAME.into(),
2615            raw_input: r#"{"command": "sleep 1000", "cd": ".", "timeout_ms": 100}"#.into(),
2616            input: json!({"command": "sleep 1000", "cd": ".", "timeout_ms": 100}),
2617            is_input_complete: true,
2618            thought_signature: None,
2619        },
2620    ));
2621    fake_model.end_last_completion_stream();
2622
2623    // Wait for the terminal tool to start running
2624    wait_for_terminal_tool_started(&mut events, cx).await;
2625
2626    // Advance clock past the timeout
2627    cx.executor().advance_clock(Duration::from_millis(200));
2628    cx.run_until_parked();
2629
2630    // The thread continues after tool completion - simulate the model ending its turn
2631    fake_model
2632        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2633    fake_model.end_last_completion_stream();
2634
2635    // Collect remaining events
2636    let remaining_events = collect_events_until_stop(&mut events, cx).await;
2637
2638    // Verify the terminal was killed due to timeout
2639    assert!(
2640        handle.was_killed(),
2641        "expected terminal handle to be killed on timeout"
2642    );
2643
2644    // Verify we got an EndTurn (the tool completed, just with timeout)
2645    assert_eq!(
2646        stop_events(remaining_events),
2647        vec![acp::StopReason::EndTurn],
2648    );
2649
2650    // Verify the tool result indicates timeout, not user stopped
2651    thread.update(cx, |thread, _cx| {
2652        let message = thread.last_received_or_pending_message().unwrap();
2653        let agent_message = message.as_agent_message().unwrap();
2654
2655        let tool_use = agent_message
2656            .content
2657            .iter()
2658            .find_map(|content| match content {
2659                AgentMessageContent::ToolUse(tool_use) => Some(tool_use),
2660                _ => None,
2661            })
2662            .expect("expected tool use in agent message");
2663
2664        let tool_result = agent_message
2665            .tool_results
2666            .get(&tool_use.id)
2667            .expect("expected tool result");
2668
2669        let result_text = match &tool_result.content {
2670            language_model::LanguageModelToolResultContent::Text(text) => text.to_string(),
2671            _ => panic!("expected text content in tool result"),
2672        };
2673
2674        assert!(
2675            result_text.contains("timed out"),
2676            "expected tool result to indicate timeout, got: {result_text}"
2677        );
2678        assert!(
2679            !result_text.contains("The user stopped"),
2680            "tool result should not mention user stopped when it timed out, got: {result_text}"
2681        );
2682    });
2683}
2684
2685#[gpui::test]
2686async fn test_in_progress_send_canceled_by_next_send(cx: &mut TestAppContext) {
2687    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2688    let fake_model = model.as_fake();
2689
2690    let events_1 = thread
2691        .update(cx, |thread, cx| {
2692            thread.send(UserMessageId::new(), ["Hello 1"], cx)
2693        })
2694        .unwrap();
2695    cx.run_until_parked();
2696    fake_model.send_last_completion_stream_text_chunk("Hey 1!");
2697    cx.run_until_parked();
2698
2699    let events_2 = thread
2700        .update(cx, |thread, cx| {
2701            thread.send(UserMessageId::new(), ["Hello 2"], cx)
2702        })
2703        .unwrap();
2704    cx.run_until_parked();
2705    fake_model.send_last_completion_stream_text_chunk("Hey 2!");
2706    fake_model
2707        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2708    fake_model.end_last_completion_stream();
2709
2710    let events_1 = events_1.collect::<Vec<_>>().await;
2711    assert_eq!(stop_events(events_1), vec![acp::StopReason::Cancelled]);
2712    let events_2 = events_2.collect::<Vec<_>>().await;
2713    assert_eq!(stop_events(events_2), vec![acp::StopReason::EndTurn]);
2714}
2715
2716#[gpui::test]
2717async fn test_retry_cancelled_promptly_on_new_send(cx: &mut TestAppContext) {
2718    // Regression test: when a completion fails with a retryable error (e.g. upstream 500),
2719    // the retry loop waits on a timer. If the user switches models and sends a new message
2720    // during that delay, the old turn should exit immediately instead of retrying with the
2721    // stale model.
2722    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2723    let model_a = model.as_fake();
2724
2725    // Start a turn with model_a.
2726    let events_1 = thread
2727        .update(cx, |thread, cx| {
2728            thread.send(UserMessageId::new(), ["Hello"], cx)
2729        })
2730        .unwrap();
2731    cx.run_until_parked();
2732    assert_eq!(model_a.completion_count(), 1);
2733
2734    // Model returns a retryable upstream 500. The turn enters the retry delay.
2735    model_a.send_last_completion_stream_error(
2736        LanguageModelCompletionError::UpstreamProviderError {
2737            message: "Internal server error".to_string(),
2738            status: http_client::StatusCode::INTERNAL_SERVER_ERROR,
2739            retry_after: None,
2740        },
2741    );
2742    model_a.end_last_completion_stream();
2743    cx.run_until_parked();
2744
2745    // The old completion was consumed; model_a has no pending requests yet because the
2746    // retry timer hasn't fired.
2747    assert_eq!(model_a.completion_count(), 0);
2748
2749    // Switch to model_b and send a new message. This cancels the old turn.
2750    let model_b = Arc::new(FakeLanguageModel::with_id_and_thinking(
2751        "fake", "model-b", "Model B", false,
2752    ));
2753    thread.update(cx, |thread, cx| {
2754        thread.set_model(model_b.clone(), cx);
2755    });
2756    let events_2 = thread
2757        .update(cx, |thread, cx| {
2758            thread.send(UserMessageId::new(), ["Continue"], cx)
2759        })
2760        .unwrap();
2761    cx.run_until_parked();
2762
2763    // model_b should have received its completion request.
2764    assert_eq!(model_b.as_fake().completion_count(), 1);
2765
2766    // Advance the clock well past the retry delay (BASE_RETRY_DELAY = 5s).
2767    cx.executor().advance_clock(Duration::from_secs(10));
2768    cx.run_until_parked();
2769
2770    // model_a must NOT have received another completion request — the cancelled turn
2771    // should have exited during the retry delay rather than retrying with the old model.
2772    assert_eq!(
2773        model_a.completion_count(),
2774        0,
2775        "old model should not receive a retry request after cancellation"
2776    );
2777
2778    // Complete model_b's turn.
2779    model_b
2780        .as_fake()
2781        .send_last_completion_stream_text_chunk("Done!");
2782    model_b
2783        .as_fake()
2784        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2785    model_b.as_fake().end_last_completion_stream();
2786
2787    let events_1 = events_1.collect::<Vec<_>>().await;
2788    assert_eq!(stop_events(events_1), vec![acp::StopReason::Cancelled]);
2789
2790    let events_2 = events_2.collect::<Vec<_>>().await;
2791    assert_eq!(stop_events(events_2), vec![acp::StopReason::EndTurn]);
2792}
2793
2794#[gpui::test]
2795async fn test_subsequent_successful_sends_dont_cancel(cx: &mut TestAppContext) {
2796    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2797    let fake_model = model.as_fake();
2798
2799    let events_1 = thread
2800        .update(cx, |thread, cx| {
2801            thread.send(UserMessageId::new(), ["Hello 1"], cx)
2802        })
2803        .unwrap();
2804    cx.run_until_parked();
2805    fake_model.send_last_completion_stream_text_chunk("Hey 1!");
2806    fake_model
2807        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2808    fake_model.end_last_completion_stream();
2809    let events_1 = events_1.collect::<Vec<_>>().await;
2810
2811    let events_2 = thread
2812        .update(cx, |thread, cx| {
2813            thread.send(UserMessageId::new(), ["Hello 2"], cx)
2814        })
2815        .unwrap();
2816    cx.run_until_parked();
2817    fake_model.send_last_completion_stream_text_chunk("Hey 2!");
2818    fake_model
2819        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn));
2820    fake_model.end_last_completion_stream();
2821    let events_2 = events_2.collect::<Vec<_>>().await;
2822
2823    assert_eq!(stop_events(events_1), vec![acp::StopReason::EndTurn]);
2824    assert_eq!(stop_events(events_2), vec![acp::StopReason::EndTurn]);
2825}
2826
2827#[gpui::test]
2828async fn test_refusal(cx: &mut TestAppContext) {
2829    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2830    let fake_model = model.as_fake();
2831
2832    let events = thread
2833        .update(cx, |thread, cx| {
2834            thread.send(UserMessageId::new(), ["Hello"], cx)
2835        })
2836        .unwrap();
2837    cx.run_until_parked();
2838    thread.read_with(cx, |thread, _| {
2839        assert_eq!(
2840            thread.to_markdown(),
2841            indoc! {"
2842                ## User
2843
2844                Hello
2845            "}
2846        );
2847    });
2848
2849    fake_model.send_last_completion_stream_text_chunk("Hey!");
2850    cx.run_until_parked();
2851    thread.read_with(cx, |thread, _| {
2852        assert_eq!(
2853            thread.to_markdown(),
2854            indoc! {"
2855                ## User
2856
2857                Hello
2858
2859                ## Assistant
2860
2861                Hey!
2862            "}
2863        );
2864    });
2865
2866    // If the model refuses to continue, the thread should remove all the messages after the last user message.
2867    fake_model
2868        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::Refusal));
2869    let events = events.collect::<Vec<_>>().await;
2870    assert_eq!(stop_events(events), vec![acp::StopReason::Refusal]);
2871    thread.read_with(cx, |thread, _| {
2872        assert_eq!(thread.to_markdown(), "");
2873    });
2874}
2875
2876#[gpui::test]
2877async fn test_truncate_first_message(cx: &mut TestAppContext) {
2878    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
2879    let fake_model = model.as_fake();
2880
2881    let message_id = UserMessageId::new();
2882    thread
2883        .update(cx, |thread, cx| {
2884            thread.send(message_id.clone(), ["Hello"], cx)
2885        })
2886        .unwrap();
2887    cx.run_until_parked();
2888    thread.read_with(cx, |thread, _| {
2889        assert_eq!(
2890            thread.to_markdown(),
2891            indoc! {"
2892                ## User
2893
2894                Hello
2895            "}
2896        );
2897        assert_eq!(thread.latest_token_usage(), None);
2898    });
2899
2900    fake_model.send_last_completion_stream_text_chunk("Hey!");
2901    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
2902        language_model::TokenUsage {
2903            input_tokens: 32_000,
2904            output_tokens: 16_000,
2905            cache_creation_input_tokens: 0,
2906            cache_read_input_tokens: 0,
2907        },
2908    ));
2909    cx.run_until_parked();
2910    thread.read_with(cx, |thread, _| {
2911        assert_eq!(
2912            thread.to_markdown(),
2913            indoc! {"
2914                ## User
2915
2916                Hello
2917
2918                ## Assistant
2919
2920                Hey!
2921            "}
2922        );
2923        assert_eq!(
2924            thread.latest_token_usage(),
2925            Some(acp_thread::TokenUsage {
2926                used_tokens: 32_000 + 16_000,
2927                max_tokens: 1_000_000,
2928                max_output_tokens: None,
2929                input_tokens: 32_000,
2930                output_tokens: 16_000,
2931            })
2932        );
2933    });
2934
2935    thread
2936        .update(cx, |thread, cx| thread.truncate(message_id, cx))
2937        .unwrap();
2938    cx.run_until_parked();
2939    thread.read_with(cx, |thread, _| {
2940        assert_eq!(thread.to_markdown(), "");
2941        assert_eq!(thread.latest_token_usage(), None);
2942    });
2943
2944    // Ensure we can still send a new message after truncation.
2945    thread
2946        .update(cx, |thread, cx| {
2947            thread.send(UserMessageId::new(), ["Hi"], cx)
2948        })
2949        .unwrap();
2950    thread.update(cx, |thread, _cx| {
2951        assert_eq!(
2952            thread.to_markdown(),
2953            indoc! {"
2954                ## User
2955
2956                Hi
2957            "}
2958        );
2959    });
2960    cx.run_until_parked();
2961    fake_model.send_last_completion_stream_text_chunk("Ahoy!");
2962    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
2963        language_model::TokenUsage {
2964            input_tokens: 40_000,
2965            output_tokens: 20_000,
2966            cache_creation_input_tokens: 0,
2967            cache_read_input_tokens: 0,
2968        },
2969    ));
2970    cx.run_until_parked();
2971    thread.read_with(cx, |thread, _| {
2972        assert_eq!(
2973            thread.to_markdown(),
2974            indoc! {"
2975                ## User
2976
2977                Hi
2978
2979                ## Assistant
2980
2981                Ahoy!
2982            "}
2983        );
2984
2985        assert_eq!(
2986            thread.latest_token_usage(),
2987            Some(acp_thread::TokenUsage {
2988                used_tokens: 40_000 + 20_000,
2989                max_tokens: 1_000_000,
2990                max_output_tokens: None,
2991                input_tokens: 40_000,
2992                output_tokens: 20_000,
2993            })
2994        );
2995    });
2996}
2997
2998#[gpui::test]
2999async fn test_truncate_second_message(cx: &mut TestAppContext) {
3000    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
3001    let fake_model = model.as_fake();
3002
3003    thread
3004        .update(cx, |thread, cx| {
3005            thread.send(UserMessageId::new(), ["Message 1"], cx)
3006        })
3007        .unwrap();
3008    cx.run_until_parked();
3009    fake_model.send_last_completion_stream_text_chunk("Message 1 response");
3010    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
3011        language_model::TokenUsage {
3012            input_tokens: 32_000,
3013            output_tokens: 16_000,
3014            cache_creation_input_tokens: 0,
3015            cache_read_input_tokens: 0,
3016        },
3017    ));
3018    fake_model.end_last_completion_stream();
3019    cx.run_until_parked();
3020
3021    let assert_first_message_state = |cx: &mut TestAppContext| {
3022        thread.clone().read_with(cx, |thread, _| {
3023            assert_eq!(
3024                thread.to_markdown(),
3025                indoc! {"
3026                    ## User
3027
3028                    Message 1
3029
3030                    ## Assistant
3031
3032                    Message 1 response
3033                "}
3034            );
3035
3036            assert_eq!(
3037                thread.latest_token_usage(),
3038                Some(acp_thread::TokenUsage {
3039                    used_tokens: 32_000 + 16_000,
3040                    max_tokens: 1_000_000,
3041                    max_output_tokens: None,
3042                    input_tokens: 32_000,
3043                    output_tokens: 16_000,
3044                })
3045            );
3046        });
3047    };
3048
3049    assert_first_message_state(cx);
3050
3051    let second_message_id = UserMessageId::new();
3052    thread
3053        .update(cx, |thread, cx| {
3054            thread.send(second_message_id.clone(), ["Message 2"], cx)
3055        })
3056        .unwrap();
3057    cx.run_until_parked();
3058
3059    fake_model.send_last_completion_stream_text_chunk("Message 2 response");
3060    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
3061        language_model::TokenUsage {
3062            input_tokens: 40_000,
3063            output_tokens: 20_000,
3064            cache_creation_input_tokens: 0,
3065            cache_read_input_tokens: 0,
3066        },
3067    ));
3068    fake_model.end_last_completion_stream();
3069    cx.run_until_parked();
3070
3071    thread.read_with(cx, |thread, _| {
3072        assert_eq!(
3073            thread.to_markdown(),
3074            indoc! {"
3075                ## User
3076
3077                Message 1
3078
3079                ## Assistant
3080
3081                Message 1 response
3082
3083                ## User
3084
3085                Message 2
3086
3087                ## Assistant
3088
3089                Message 2 response
3090            "}
3091        );
3092
3093        assert_eq!(
3094            thread.latest_token_usage(),
3095            Some(acp_thread::TokenUsage {
3096                used_tokens: 40_000 + 20_000,
3097                max_tokens: 1_000_000,
3098                max_output_tokens: None,
3099                input_tokens: 40_000,
3100                output_tokens: 20_000,
3101            })
3102        );
3103    });
3104
3105    thread
3106        .update(cx, |thread, cx| thread.truncate(second_message_id, cx))
3107        .unwrap();
3108    cx.run_until_parked();
3109
3110    assert_first_message_state(cx);
3111}
3112
3113#[gpui::test]
3114async fn test_title_generation(cx: &mut TestAppContext) {
3115    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
3116    let fake_model = model.as_fake();
3117
3118    let summary_model = Arc::new(FakeLanguageModel::default());
3119    thread.update(cx, |thread, cx| {
3120        thread.set_summarization_model(Some(summary_model.clone()), cx)
3121    });
3122
3123    let send = thread
3124        .update(cx, |thread, cx| {
3125            thread.send(UserMessageId::new(), ["Hello"], cx)
3126        })
3127        .unwrap();
3128    cx.run_until_parked();
3129
3130    fake_model.send_last_completion_stream_text_chunk("Hey!");
3131    fake_model.end_last_completion_stream();
3132    cx.run_until_parked();
3133    thread.read_with(cx, |thread, _| assert_eq!(thread.title(), None));
3134
3135    // Ensure the summary model has been invoked to generate a title.
3136    summary_model.send_last_completion_stream_text_chunk("Hello ");
3137    summary_model.send_last_completion_stream_text_chunk("world\nG");
3138    summary_model.send_last_completion_stream_text_chunk("oodnight Moon");
3139    summary_model.end_last_completion_stream();
3140    send.collect::<Vec<_>>().await;
3141    cx.run_until_parked();
3142    thread.read_with(cx, |thread, _| {
3143        assert_eq!(thread.title(), Some("Hello world".into()))
3144    });
3145
3146    // Send another message, ensuring no title is generated this time.
3147    let send = thread
3148        .update(cx, |thread, cx| {
3149            thread.send(UserMessageId::new(), ["Hello again"], cx)
3150        })
3151        .unwrap();
3152    cx.run_until_parked();
3153    fake_model.send_last_completion_stream_text_chunk("Hey again!");
3154    fake_model.end_last_completion_stream();
3155    cx.run_until_parked();
3156    assert_eq!(summary_model.pending_completions(), Vec::new());
3157    send.collect::<Vec<_>>().await;
3158    thread.read_with(cx, |thread, _| {
3159        assert_eq!(thread.title(), Some("Hello world".into()))
3160    });
3161}
3162
3163#[gpui::test]
3164async fn test_building_request_with_pending_tools(cx: &mut TestAppContext) {
3165    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
3166    let fake_model = model.as_fake();
3167
3168    let _events = thread
3169        .update(cx, |thread, cx| {
3170            thread.add_tool(ToolRequiringPermission);
3171            thread.add_tool(EchoTool);
3172            thread.send(UserMessageId::new(), ["Hey!"], cx)
3173        })
3174        .unwrap();
3175    cx.run_until_parked();
3176
3177    let permission_tool_use = LanguageModelToolUse {
3178        id: "tool_id_1".into(),
3179        name: ToolRequiringPermission::NAME.into(),
3180        raw_input: "{}".into(),
3181        input: json!({}),
3182        is_input_complete: true,
3183        thought_signature: None,
3184    };
3185    let echo_tool_use = LanguageModelToolUse {
3186        id: "tool_id_2".into(),
3187        name: EchoTool::NAME.into(),
3188        raw_input: json!({"text": "test"}).to_string(),
3189        input: json!({"text": "test"}),
3190        is_input_complete: true,
3191        thought_signature: None,
3192    };
3193    fake_model.send_last_completion_stream_text_chunk("Hi!");
3194    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3195        permission_tool_use,
3196    ));
3197    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3198        echo_tool_use.clone(),
3199    ));
3200    fake_model.end_last_completion_stream();
3201    cx.run_until_parked();
3202
3203    // Ensure pending tools are skipped when building a request.
3204    let request = thread
3205        .read_with(cx, |thread, cx| {
3206            thread.build_completion_request(CompletionIntent::EditFile, cx)
3207        })
3208        .unwrap();
3209    assert_eq!(
3210        request.messages[1..],
3211        vec![
3212            LanguageModelRequestMessage {
3213                role: Role::User,
3214                content: vec!["Hey!".into()],
3215                cache: true,
3216                reasoning_details: None,
3217            },
3218            LanguageModelRequestMessage {
3219                role: Role::Assistant,
3220                content: vec![
3221                    MessageContent::Text("Hi!".into()),
3222                    MessageContent::ToolUse(echo_tool_use.clone())
3223                ],
3224                cache: false,
3225                reasoning_details: None,
3226            },
3227            LanguageModelRequestMessage {
3228                role: Role::User,
3229                content: vec![MessageContent::ToolResult(LanguageModelToolResult {
3230                    tool_use_id: echo_tool_use.id.clone(),
3231                    tool_name: echo_tool_use.name,
3232                    is_error: false,
3233                    content: "test".into(),
3234                    output: Some("test".into())
3235                })],
3236                cache: false,
3237                reasoning_details: None,
3238            },
3239        ],
3240    );
3241}
3242
3243#[gpui::test]
3244async fn test_agent_connection(cx: &mut TestAppContext) {
3245    cx.update(settings::init);
3246    let templates = Templates::new();
3247
3248    // Initialize language model system with test provider
3249    cx.update(|cx| {
3250        gpui_tokio::init(cx);
3251
3252        let http_client = FakeHttpClient::with_404_response();
3253        let clock = Arc::new(clock::FakeSystemClock::new());
3254        let client = Client::new(clock, http_client, cx);
3255        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
3256        language_model::init(cx);
3257        RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
3258        language_models::init(user_store, client.clone(), cx);
3259        LanguageModelRegistry::test(cx);
3260    });
3261    cx.executor().forbid_parking();
3262
3263    // Create a project for new_thread
3264    let fake_fs = cx.update(|cx| fs::FakeFs::new(cx.background_executor().clone()));
3265    fake_fs.insert_tree(path!("/test"), json!({})).await;
3266    let project = Project::test(fake_fs.clone(), [Path::new("/test")], cx).await;
3267    let cwd = PathList::new(&[Path::new("/test")]);
3268    let thread_store = cx.new(|cx| ThreadStore::new(cx));
3269
3270    // Create agent and connection
3271    let agent = cx
3272        .update(|cx| NativeAgent::new(thread_store, templates.clone(), None, fake_fs.clone(), cx));
3273    let connection = NativeAgentConnection(agent.clone());
3274
3275    // Create a thread using new_thread
3276    let connection_rc = Rc::new(connection.clone());
3277    let acp_thread = cx
3278        .update(|cx| connection_rc.new_session(project, cwd, cx))
3279        .await
3280        .expect("new_thread should succeed");
3281
3282    // Get the session_id from the AcpThread
3283    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
3284
3285    // Test model_selector returns Some
3286    let selector_opt = connection.model_selector(&session_id);
3287    assert!(
3288        selector_opt.is_some(),
3289        "agent should always support ModelSelector"
3290    );
3291    let selector = selector_opt.unwrap();
3292
3293    // Test list_models
3294    let listed_models = cx
3295        .update(|cx| selector.list_models(cx))
3296        .await
3297        .expect("list_models should succeed");
3298    let AgentModelList::Grouped(listed_models) = listed_models else {
3299        panic!("Unexpected model list type");
3300    };
3301    assert!(!listed_models.is_empty(), "should have at least one model");
3302    assert_eq!(
3303        listed_models[&AgentModelGroupName("Fake".into())][0]
3304            .id
3305            .0
3306            .as_ref(),
3307        "fake/fake"
3308    );
3309
3310    // Test selected_model returns the default
3311    let model = cx
3312        .update(|cx| selector.selected_model(cx))
3313        .await
3314        .expect("selected_model should succeed");
3315    let model = cx
3316        .update(|cx| agent.read(cx).models().model_from_id(&model.id))
3317        .unwrap();
3318    let model = model.as_fake();
3319    assert_eq!(model.id().0, "fake", "should return default model");
3320
3321    let request = acp_thread.update(cx, |thread, cx| thread.send(vec!["abc".into()], cx));
3322    cx.run_until_parked();
3323    model.send_last_completion_stream_text_chunk("def");
3324    cx.run_until_parked();
3325    acp_thread.read_with(cx, |thread, cx| {
3326        assert_eq!(
3327            thread.to_markdown(cx),
3328            indoc! {"
3329                ## User
3330
3331                abc
3332
3333                ## Assistant
3334
3335                def
3336
3337            "}
3338        )
3339    });
3340
3341    // Test cancel
3342    cx.update(|cx| connection.cancel(&session_id, cx));
3343    request.await.expect("prompt should fail gracefully");
3344
3345    // Explicitly close the session and drop the ACP thread.
3346    cx.update(|cx| Rc::new(connection.clone()).close_session(&session_id, cx))
3347        .await
3348        .unwrap();
3349    drop(acp_thread);
3350    let result = cx
3351        .update(|cx| {
3352            connection.prompt(
3353                Some(acp_thread::UserMessageId::new()),
3354                acp::PromptRequest::new(session_id.clone(), vec!["ghi".into()]),
3355                cx,
3356            )
3357        })
3358        .await;
3359    assert_eq!(
3360        result.as_ref().unwrap_err().to_string(),
3361        "Session not found",
3362        "unexpected result: {:?}",
3363        result
3364    );
3365}
3366
3367#[gpui::test]
3368async fn test_tool_updates_to_completion(cx: &mut TestAppContext) {
3369    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3370    thread.update(cx, |thread, _cx| thread.add_tool(EchoTool));
3371    let fake_model = model.as_fake();
3372
3373    let mut events = thread
3374        .update(cx, |thread, cx| {
3375            thread.send(UserMessageId::new(), ["Echo something"], cx)
3376        })
3377        .unwrap();
3378    cx.run_until_parked();
3379
3380    // Simulate streaming partial input.
3381    let input = json!({});
3382    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3383        LanguageModelToolUse {
3384            id: "1".into(),
3385            name: EchoTool::NAME.into(),
3386            raw_input: input.to_string(),
3387            input,
3388            is_input_complete: false,
3389            thought_signature: None,
3390        },
3391    ));
3392
3393    // Input streaming completed
3394    let input = json!({ "text": "Hello!" });
3395    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3396        LanguageModelToolUse {
3397            id: "1".into(),
3398            name: "echo".into(),
3399            raw_input: input.to_string(),
3400            input,
3401            is_input_complete: true,
3402            thought_signature: None,
3403        },
3404    ));
3405    fake_model.end_last_completion_stream();
3406    cx.run_until_parked();
3407
3408    let tool_call = expect_tool_call(&mut events).await;
3409    assert_eq!(
3410        tool_call,
3411        acp::ToolCall::new("1", "Echo")
3412            .raw_input(json!({}))
3413            .meta(acp::Meta::from_iter([("tool_name".into(), "echo".into())]))
3414    );
3415    let update = expect_tool_call_update_fields(&mut events).await;
3416    assert_eq!(
3417        update,
3418        acp::ToolCallUpdate::new(
3419            "1",
3420            acp::ToolCallUpdateFields::new()
3421                .title("Echo")
3422                .kind(acp::ToolKind::Other)
3423                .raw_input(json!({ "text": "Hello!"}))
3424        )
3425    );
3426    let update = expect_tool_call_update_fields(&mut events).await;
3427    assert_eq!(
3428        update,
3429        acp::ToolCallUpdate::new(
3430            "1",
3431            acp::ToolCallUpdateFields::new().status(acp::ToolCallStatus::InProgress)
3432        )
3433    );
3434    let update = expect_tool_call_update_fields(&mut events).await;
3435    assert_eq!(
3436        update,
3437        acp::ToolCallUpdate::new(
3438            "1",
3439            acp::ToolCallUpdateFields::new()
3440                .status(acp::ToolCallStatus::Completed)
3441                .raw_output("Hello!")
3442        )
3443    );
3444}
3445
3446#[gpui::test]
3447async fn test_update_plan_tool_updates_thread_events(cx: &mut TestAppContext) {
3448    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3449    thread.update(cx, |thread, _cx| thread.add_tool(UpdatePlanTool));
3450    let fake_model = model.as_fake();
3451
3452    let mut events = thread
3453        .update(cx, |thread, cx| {
3454            thread.send(UserMessageId::new(), ["Make a plan"], cx)
3455        })
3456        .unwrap();
3457    cx.run_until_parked();
3458
3459    let input = json!({
3460        "plan": [
3461            {
3462                "step": "Inspect the code",
3463                "status": "completed",
3464            },
3465            {
3466                "step": "Implement the tool",
3467                "status": "in_progress"
3468            },
3469            {
3470                "step": "Run tests",
3471                "status": "pending",
3472            }
3473        ]
3474    });
3475    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3476        LanguageModelToolUse {
3477            id: "plan_1".into(),
3478            name: UpdatePlanTool::NAME.into(),
3479            raw_input: input.to_string(),
3480            input,
3481            is_input_complete: true,
3482            thought_signature: None,
3483        },
3484    ));
3485    fake_model.end_last_completion_stream();
3486    cx.run_until_parked();
3487
3488    let tool_call = expect_tool_call(&mut events).await;
3489    assert_eq!(
3490        tool_call,
3491        acp::ToolCall::new("plan_1", "Update plan")
3492            .kind(acp::ToolKind::Think)
3493            .raw_input(json!({
3494                "plan": [
3495                    {
3496                        "step": "Inspect the code",
3497                        "status": "completed",
3498                    },
3499                    {
3500                        "step": "Implement the tool",
3501                        "status": "in_progress"
3502                    },
3503                    {
3504                        "step": "Run tests",
3505                        "status": "pending",
3506                    }
3507                ]
3508            }))
3509            .meta(acp::Meta::from_iter([(
3510                "tool_name".into(),
3511                "update_plan".into()
3512            )]))
3513    );
3514
3515    let update = expect_tool_call_update_fields(&mut events).await;
3516    assert_eq!(
3517        update,
3518        acp::ToolCallUpdate::new(
3519            "plan_1",
3520            acp::ToolCallUpdateFields::new().status(acp::ToolCallStatus::InProgress)
3521        )
3522    );
3523
3524    let plan = expect_plan(&mut events).await;
3525    assert_eq!(
3526        plan,
3527        acp::Plan::new(vec![
3528            acp::PlanEntry::new(
3529                "Inspect the code",
3530                acp::PlanEntryPriority::Medium,
3531                acp::PlanEntryStatus::Completed,
3532            ),
3533            acp::PlanEntry::new(
3534                "Implement the tool",
3535                acp::PlanEntryPriority::Medium,
3536                acp::PlanEntryStatus::InProgress,
3537            ),
3538            acp::PlanEntry::new(
3539                "Run tests",
3540                acp::PlanEntryPriority::Medium,
3541                acp::PlanEntryStatus::Pending,
3542            ),
3543        ])
3544    );
3545
3546    let update = expect_tool_call_update_fields(&mut events).await;
3547    assert_eq!(
3548        update,
3549        acp::ToolCallUpdate::new(
3550            "plan_1",
3551            acp::ToolCallUpdateFields::new()
3552                .status(acp::ToolCallStatus::Completed)
3553                .raw_output("Plan updated")
3554        )
3555    );
3556}
3557
3558#[gpui::test]
3559async fn test_send_no_retry_on_success(cx: &mut TestAppContext) {
3560    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3561    let fake_model = model.as_fake();
3562
3563    let mut events = thread
3564        .update(cx, |thread, cx| {
3565            thread.send(UserMessageId::new(), ["Hello!"], cx)
3566        })
3567        .unwrap();
3568    cx.run_until_parked();
3569
3570    fake_model.send_last_completion_stream_text_chunk("Hey!");
3571    fake_model.end_last_completion_stream();
3572
3573    let mut retry_events = Vec::new();
3574    while let Some(Ok(event)) = events.next().await {
3575        match event {
3576            ThreadEvent::Retry(retry_status) => {
3577                retry_events.push(retry_status);
3578            }
3579            ThreadEvent::Stop(..) => break,
3580            _ => {}
3581        }
3582    }
3583
3584    assert_eq!(retry_events.len(), 0);
3585    thread.read_with(cx, |thread, _cx| {
3586        assert_eq!(
3587            thread.to_markdown(),
3588            indoc! {"
3589                ## User
3590
3591                Hello!
3592
3593                ## Assistant
3594
3595                Hey!
3596            "}
3597        )
3598    });
3599}
3600
3601#[gpui::test]
3602async fn test_send_retry_on_error(cx: &mut TestAppContext) {
3603    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3604    let fake_model = model.as_fake();
3605
3606    let mut events = thread
3607        .update(cx, |thread, cx| {
3608            thread.send(UserMessageId::new(), ["Hello!"], cx)
3609        })
3610        .unwrap();
3611    cx.run_until_parked();
3612
3613    fake_model.send_last_completion_stream_text_chunk("Hey,");
3614    fake_model.send_last_completion_stream_error(LanguageModelCompletionError::ServerOverloaded {
3615        provider: LanguageModelProviderName::new("Anthropic"),
3616        retry_after: Some(Duration::from_secs(3)),
3617    });
3618    fake_model.end_last_completion_stream();
3619
3620    cx.executor().advance_clock(Duration::from_secs(3));
3621    cx.run_until_parked();
3622
3623    fake_model.send_last_completion_stream_text_chunk("there!");
3624    fake_model.end_last_completion_stream();
3625    cx.run_until_parked();
3626
3627    let mut retry_events = Vec::new();
3628    while let Some(Ok(event)) = events.next().await {
3629        match event {
3630            ThreadEvent::Retry(retry_status) => {
3631                retry_events.push(retry_status);
3632            }
3633            ThreadEvent::Stop(..) => break,
3634            _ => {}
3635        }
3636    }
3637
3638    assert_eq!(retry_events.len(), 1);
3639    assert!(matches!(
3640        retry_events[0],
3641        acp_thread::RetryStatus { attempt: 1, .. }
3642    ));
3643    thread.read_with(cx, |thread, _cx| {
3644        assert_eq!(
3645            thread.to_markdown(),
3646            indoc! {"
3647                ## User
3648
3649                Hello!
3650
3651                ## Assistant
3652
3653                Hey,
3654
3655                [resume]
3656
3657                ## Assistant
3658
3659                there!
3660            "}
3661        )
3662    });
3663}
3664
3665#[gpui::test]
3666async fn test_send_retry_finishes_tool_calls_on_error(cx: &mut TestAppContext) {
3667    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3668    let fake_model = model.as_fake();
3669
3670    let events = thread
3671        .update(cx, |thread, cx| {
3672            thread.add_tool(EchoTool);
3673            thread.send(UserMessageId::new(), ["Call the echo tool!"], cx)
3674        })
3675        .unwrap();
3676    cx.run_until_parked();
3677
3678    let tool_use_1 = LanguageModelToolUse {
3679        id: "tool_1".into(),
3680        name: EchoTool::NAME.into(),
3681        raw_input: json!({"text": "test"}).to_string(),
3682        input: json!({"text": "test"}),
3683        is_input_complete: true,
3684        thought_signature: None,
3685    };
3686    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
3687        tool_use_1.clone(),
3688    ));
3689    fake_model.send_last_completion_stream_error(LanguageModelCompletionError::ServerOverloaded {
3690        provider: LanguageModelProviderName::new("Anthropic"),
3691        retry_after: Some(Duration::from_secs(3)),
3692    });
3693    fake_model.end_last_completion_stream();
3694
3695    cx.executor().advance_clock(Duration::from_secs(3));
3696    let completion = fake_model.pending_completions().pop().unwrap();
3697    assert_eq!(
3698        completion.messages[1..],
3699        vec![
3700            LanguageModelRequestMessage {
3701                role: Role::User,
3702                content: vec!["Call the echo tool!".into()],
3703                cache: false,
3704                reasoning_details: None,
3705            },
3706            LanguageModelRequestMessage {
3707                role: Role::Assistant,
3708                content: vec![language_model::MessageContent::ToolUse(tool_use_1.clone())],
3709                cache: false,
3710                reasoning_details: None,
3711            },
3712            LanguageModelRequestMessage {
3713                role: Role::User,
3714                content: vec![language_model::MessageContent::ToolResult(
3715                    LanguageModelToolResult {
3716                        tool_use_id: tool_use_1.id.clone(),
3717                        tool_name: tool_use_1.name.clone(),
3718                        is_error: false,
3719                        content: "test".into(),
3720                        output: Some("test".into())
3721                    }
3722                )],
3723                cache: true,
3724                reasoning_details: None,
3725            },
3726        ]
3727    );
3728
3729    fake_model.send_last_completion_stream_text_chunk("Done");
3730    fake_model.end_last_completion_stream();
3731    cx.run_until_parked();
3732    events.collect::<Vec<_>>().await;
3733    thread.read_with(cx, |thread, _cx| {
3734        assert_eq!(
3735            thread.last_received_or_pending_message(),
3736            Some(Message::Agent(AgentMessage {
3737                content: vec![AgentMessageContent::Text("Done".into())],
3738                tool_results: IndexMap::default(),
3739                reasoning_details: None,
3740            }))
3741        );
3742    })
3743}
3744
3745#[gpui::test]
3746async fn test_send_max_retries_exceeded(cx: &mut TestAppContext) {
3747    let ThreadTest { thread, model, .. } = setup(cx, TestModel::Fake).await;
3748    let fake_model = model.as_fake();
3749
3750    let mut events = thread
3751        .update(cx, |thread, cx| {
3752            thread.send(UserMessageId::new(), ["Hello!"], cx)
3753        })
3754        .unwrap();
3755    cx.run_until_parked();
3756
3757    for _ in 0..crate::thread::MAX_RETRY_ATTEMPTS + 1 {
3758        fake_model.send_last_completion_stream_error(
3759            LanguageModelCompletionError::ServerOverloaded {
3760                provider: LanguageModelProviderName::new("Anthropic"),
3761                retry_after: Some(Duration::from_secs(3)),
3762            },
3763        );
3764        fake_model.end_last_completion_stream();
3765        cx.executor().advance_clock(Duration::from_secs(3));
3766        cx.run_until_parked();
3767    }
3768
3769    let mut errors = Vec::new();
3770    let mut retry_events = Vec::new();
3771    while let Some(event) = events.next().await {
3772        match event {
3773            Ok(ThreadEvent::Retry(retry_status)) => {
3774                retry_events.push(retry_status);
3775            }
3776            Ok(ThreadEvent::Stop(..)) => break,
3777            Err(error) => errors.push(error),
3778            _ => {}
3779        }
3780    }
3781
3782    assert_eq!(
3783        retry_events.len(),
3784        crate::thread::MAX_RETRY_ATTEMPTS as usize
3785    );
3786    for i in 0..crate::thread::MAX_RETRY_ATTEMPTS as usize {
3787        assert_eq!(retry_events[i].attempt, i + 1);
3788    }
3789    assert_eq!(errors.len(), 1);
3790    let error = errors[0]
3791        .downcast_ref::<LanguageModelCompletionError>()
3792        .unwrap();
3793    assert!(matches!(
3794        error,
3795        LanguageModelCompletionError::ServerOverloaded { .. }
3796    ));
3797}
3798
3799#[gpui::test]
3800async fn test_streaming_tool_completes_when_llm_stream_ends_without_final_input(
3801    cx: &mut TestAppContext,
3802) {
3803    init_test(cx);
3804    always_allow_tools(cx);
3805
3806    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
3807    let fake_model = model.as_fake();
3808
3809    thread.update(cx, |thread, _cx| {
3810        thread.add_tool(StreamingEchoTool::new());
3811    });
3812
3813    let _events = thread
3814        .update(cx, |thread, cx| {
3815            thread.send(UserMessageId::new(), ["Use the streaming_echo tool"], cx)
3816        })
3817        .unwrap();
3818    cx.run_until_parked();
3819
3820    // Send a partial tool use (is_input_complete = false), simulating the LLM
3821    // streaming input for a tool.
3822    let tool_use = LanguageModelToolUse {
3823        id: "tool_1".into(),
3824        name: "streaming_echo".into(),
3825        raw_input: r#"{"text": "partial"}"#.into(),
3826        input: json!({"text": "partial"}),
3827        is_input_complete: false,
3828        thought_signature: None,
3829    };
3830    fake_model
3831        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(tool_use.clone()));
3832    cx.run_until_parked();
3833
3834    // Send a stream error WITHOUT ever sending is_input_complete = true.
3835    // Before the fix, this would deadlock: the tool waits for more partials
3836    // (or cancellation), run_turn_internal waits for the tool, and the sender
3837    // keeping the channel open lives inside RunningTurn.
3838    fake_model.send_last_completion_stream_error(
3839        LanguageModelCompletionError::UpstreamProviderError {
3840            message: "Internal server error".to_string(),
3841            status: http_client::StatusCode::INTERNAL_SERVER_ERROR,
3842            retry_after: None,
3843        },
3844    );
3845    fake_model.end_last_completion_stream();
3846
3847    // Advance past the retry delay so run_turn_internal retries.
3848    cx.executor().advance_clock(Duration::from_secs(5));
3849    cx.run_until_parked();
3850
3851    // The retry request should contain the streaming tool's error result,
3852    // proving the tool terminated and its result was forwarded.
3853    let completion = fake_model
3854        .pending_completions()
3855        .pop()
3856        .expect("No running turn");
3857    assert_eq!(
3858        completion.messages[1..],
3859        vec![
3860            LanguageModelRequestMessage {
3861                role: Role::User,
3862                content: vec!["Use the streaming_echo tool".into()],
3863                cache: false,
3864                reasoning_details: None,
3865            },
3866            LanguageModelRequestMessage {
3867                role: Role::Assistant,
3868                content: vec![language_model::MessageContent::ToolUse(tool_use.clone())],
3869                cache: false,
3870                reasoning_details: None,
3871            },
3872            LanguageModelRequestMessage {
3873                role: Role::User,
3874                content: vec![language_model::MessageContent::ToolResult(
3875                    LanguageModelToolResult {
3876                        tool_use_id: tool_use.id.clone(),
3877                        tool_name: tool_use.name,
3878                        is_error: true,
3879                        content: "Failed to receive tool input: tool input was not fully received"
3880                            .into(),
3881                        output: Some(
3882                            "Failed to receive tool input: tool input was not fully received"
3883                                .into()
3884                        ),
3885                    }
3886                )],
3887                cache: true,
3888                reasoning_details: None,
3889            },
3890        ]
3891    );
3892
3893    // Finish the retry round so the turn completes cleanly.
3894    fake_model.send_last_completion_stream_text_chunk("Done");
3895    fake_model.end_last_completion_stream();
3896    cx.run_until_parked();
3897
3898    thread.read_with(cx, |thread, _cx| {
3899        assert!(
3900            thread.is_turn_complete(),
3901            "Thread should not be stuck; the turn should have completed",
3902        );
3903    });
3904}
3905
3906/// Filters out the stop events for asserting against in tests
3907fn stop_events(result_events: Vec<Result<ThreadEvent>>) -> Vec<acp::StopReason> {
3908    result_events
3909        .into_iter()
3910        .filter_map(|event| match event.unwrap() {
3911            ThreadEvent::Stop(stop_reason) => Some(stop_reason),
3912            _ => None,
3913        })
3914        .collect()
3915}
3916
3917struct ThreadTest {
3918    model: Arc<dyn LanguageModel>,
3919    thread: Entity<Thread>,
3920    project_context: Entity<ProjectContext>,
3921    context_server_store: Entity<ContextServerStore>,
3922    fs: Arc<FakeFs>,
3923}
3924
3925enum TestModel {
3926    Sonnet4,
3927    Fake,
3928}
3929
3930impl TestModel {
3931    fn id(&self) -> LanguageModelId {
3932        match self {
3933            TestModel::Sonnet4 => LanguageModelId("claude-sonnet-4-latest".into()),
3934            TestModel::Fake => unreachable!(),
3935        }
3936    }
3937}
3938
3939async fn setup(cx: &mut TestAppContext, model: TestModel) -> ThreadTest {
3940    cx.executor().allow_parking();
3941
3942    let fs = FakeFs::new(cx.background_executor.clone());
3943    fs.create_dir(paths::settings_file().parent().unwrap())
3944        .await
3945        .unwrap();
3946    fs.insert_file(
3947        paths::settings_file(),
3948        json!({
3949            "agent": {
3950                "default_profile": "test-profile",
3951                "profiles": {
3952                    "test-profile": {
3953                        "name": "Test Profile",
3954                        "tools": {
3955                            EchoTool::NAME: true,
3956                            DelayTool::NAME: true,
3957                            WordListTool::NAME: true,
3958                            ToolRequiringPermission::NAME: true,
3959                            InfiniteTool::NAME: true,
3960                            CancellationAwareTool::NAME: true,
3961                            StreamingEchoTool::NAME: true,
3962                            StreamingFailingEchoTool::NAME: true,
3963                            TerminalTool::NAME: true,
3964                            UpdatePlanTool::NAME: true,
3965                        }
3966                    }
3967                }
3968            }
3969        })
3970        .to_string()
3971        .into_bytes(),
3972    )
3973    .await;
3974
3975    cx.update(|cx| {
3976        settings::init(cx);
3977
3978        match model {
3979            TestModel::Fake => {}
3980            TestModel::Sonnet4 => {
3981                gpui_tokio::init(cx);
3982                let http_client = ReqwestClient::user_agent("agent tests").unwrap();
3983                cx.set_http_client(Arc::new(http_client));
3984                let client = Client::production(cx);
3985                let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
3986                language_model::init(cx);
3987                RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
3988                language_models::init(user_store, client.clone(), cx);
3989            }
3990        };
3991
3992        watch_settings(fs.clone(), cx);
3993    });
3994
3995    let templates = Templates::new();
3996
3997    fs.insert_tree(path!("/test"), json!({})).await;
3998    let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
3999
4000    let model = cx
4001        .update(|cx| {
4002            if let TestModel::Fake = model {
4003                Task::ready(Arc::new(FakeLanguageModel::default()) as Arc<_>)
4004            } else {
4005                let model_id = model.id();
4006                let models = LanguageModelRegistry::read_global(cx);
4007                let model = models
4008                    .available_models(cx)
4009                    .find(|model| model.id() == model_id)
4010                    .unwrap();
4011
4012                let provider = models.provider(&model.provider_id()).unwrap();
4013                let authenticated = provider.authenticate(cx);
4014
4015                cx.spawn(async move |_cx| {
4016                    authenticated.await.unwrap();
4017                    model
4018                })
4019            }
4020        })
4021        .await;
4022
4023    let project_context = cx.new(|_cx| ProjectContext::default());
4024    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
4025    let context_server_registry =
4026        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
4027    let thread = cx.new(|cx| {
4028        Thread::new(
4029            project,
4030            project_context.clone(),
4031            context_server_registry,
4032            templates,
4033            Some(model.clone()),
4034            cx,
4035        )
4036    });
4037    ThreadTest {
4038        model,
4039        thread,
4040        project_context,
4041        context_server_store,
4042        fs,
4043    }
4044}
4045
4046#[cfg(test)]
4047#[ctor::ctor]
4048fn init_logger() {
4049    if std::env::var("RUST_LOG").is_ok() {
4050        env_logger::init();
4051    }
4052}
4053
4054fn watch_settings(fs: Arc<dyn Fs>, cx: &mut App) {
4055    let fs = fs.clone();
4056    cx.spawn({
4057        async move |cx| {
4058            let (mut new_settings_content_rx, watcher_task) = settings::watch_config_file(
4059                cx.background_executor(),
4060                fs,
4061                paths::settings_file().clone(),
4062            );
4063            let _watcher_task = watcher_task;
4064
4065            while let Some(new_settings_content) = new_settings_content_rx.next().await {
4066                cx.update(|cx| {
4067                    SettingsStore::update_global(cx, |settings, cx| {
4068                        settings.set_user_settings(&new_settings_content, cx)
4069                    })
4070                })
4071                .ok();
4072            }
4073        }
4074    })
4075    .detach();
4076}
4077
4078fn tool_names_for_completion(completion: &LanguageModelRequest) -> Vec<String> {
4079    completion
4080        .tools
4081        .iter()
4082        .map(|tool| tool.name.clone())
4083        .collect()
4084}
4085
4086fn setup_context_server(
4087    name: &'static str,
4088    tools: Vec<context_server::types::Tool>,
4089    context_server_store: &Entity<ContextServerStore>,
4090    cx: &mut TestAppContext,
4091) -> mpsc::UnboundedReceiver<(
4092    context_server::types::CallToolParams,
4093    oneshot::Sender<context_server::types::CallToolResponse>,
4094)> {
4095    cx.update(|cx| {
4096        let mut settings = ProjectSettings::get_global(cx).clone();
4097        settings.context_servers.insert(
4098            name.into(),
4099            project::project_settings::ContextServerSettings::Stdio {
4100                enabled: true,
4101                remote: false,
4102                command: ContextServerCommand {
4103                    path: "somebinary".into(),
4104                    args: Vec::new(),
4105                    env: None,
4106                    timeout: None,
4107                },
4108            },
4109        );
4110        ProjectSettings::override_global(settings, cx);
4111    });
4112
4113    let (mcp_tool_calls_tx, mcp_tool_calls_rx) = mpsc::unbounded();
4114    let fake_transport = context_server::test::create_fake_transport(name, cx.executor())
4115        .on_request::<context_server::types::requests::Initialize, _>(move |_params| async move {
4116            context_server::types::InitializeResponse {
4117                protocol_version: context_server::types::ProtocolVersion(
4118                    context_server::types::LATEST_PROTOCOL_VERSION.to_string(),
4119                ),
4120                server_info: context_server::types::Implementation {
4121                    name: name.into(),
4122                    version: "1.0.0".to_string(),
4123                },
4124                capabilities: context_server::types::ServerCapabilities {
4125                    tools: Some(context_server::types::ToolsCapabilities {
4126                        list_changed: Some(true),
4127                    }),
4128                    ..Default::default()
4129                },
4130                meta: None,
4131            }
4132        })
4133        .on_request::<context_server::types::requests::ListTools, _>(move |_params| {
4134            let tools = tools.clone();
4135            async move {
4136                context_server::types::ListToolsResponse {
4137                    tools,
4138                    next_cursor: None,
4139                    meta: None,
4140                }
4141            }
4142        })
4143        .on_request::<context_server::types::requests::CallTool, _>(move |params| {
4144            let mcp_tool_calls_tx = mcp_tool_calls_tx.clone();
4145            async move {
4146                let (response_tx, response_rx) = oneshot::channel();
4147                mcp_tool_calls_tx
4148                    .unbounded_send((params, response_tx))
4149                    .unwrap();
4150                response_rx.await.unwrap()
4151            }
4152        });
4153    context_server_store.update(cx, |store, cx| {
4154        store.start_server(
4155            Arc::new(ContextServer::new(
4156                ContextServerId(name.into()),
4157                Arc::new(fake_transport),
4158            )),
4159            cx,
4160        );
4161    });
4162    cx.run_until_parked();
4163    mcp_tool_calls_rx
4164}
4165
4166#[gpui::test]
4167async fn test_tokens_before_message(cx: &mut TestAppContext) {
4168    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
4169    let fake_model = model.as_fake();
4170
4171    // First message
4172    let message_1_id = UserMessageId::new();
4173    thread
4174        .update(cx, |thread, cx| {
4175            thread.send(message_1_id.clone(), ["First message"], cx)
4176        })
4177        .unwrap();
4178    cx.run_until_parked();
4179
4180    // Before any response, tokens_before_message should return None for first message
4181    thread.read_with(cx, |thread, _| {
4182        assert_eq!(
4183            thread.tokens_before_message(&message_1_id),
4184            None,
4185            "First message should have no tokens before it"
4186        );
4187    });
4188
4189    // Complete first message with usage
4190    fake_model.send_last_completion_stream_text_chunk("Response 1");
4191    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
4192        language_model::TokenUsage {
4193            input_tokens: 100,
4194            output_tokens: 50,
4195            cache_creation_input_tokens: 0,
4196            cache_read_input_tokens: 0,
4197        },
4198    ));
4199    fake_model.end_last_completion_stream();
4200    cx.run_until_parked();
4201
4202    // First message still has no tokens before it
4203    thread.read_with(cx, |thread, _| {
4204        assert_eq!(
4205            thread.tokens_before_message(&message_1_id),
4206            None,
4207            "First message should still have no tokens before it after response"
4208        );
4209    });
4210
4211    // Second message
4212    let message_2_id = UserMessageId::new();
4213    thread
4214        .update(cx, |thread, cx| {
4215            thread.send(message_2_id.clone(), ["Second message"], cx)
4216        })
4217        .unwrap();
4218    cx.run_until_parked();
4219
4220    // Second message should have first message's input tokens before it
4221    thread.read_with(cx, |thread, _| {
4222        assert_eq!(
4223            thread.tokens_before_message(&message_2_id),
4224            Some(100),
4225            "Second message should have 100 tokens before it (from first request)"
4226        );
4227    });
4228
4229    // Complete second message
4230    fake_model.send_last_completion_stream_text_chunk("Response 2");
4231    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
4232        language_model::TokenUsage {
4233            input_tokens: 250, // Total for this request (includes previous context)
4234            output_tokens: 75,
4235            cache_creation_input_tokens: 0,
4236            cache_read_input_tokens: 0,
4237        },
4238    ));
4239    fake_model.end_last_completion_stream();
4240    cx.run_until_parked();
4241
4242    // Third message
4243    let message_3_id = UserMessageId::new();
4244    thread
4245        .update(cx, |thread, cx| {
4246            thread.send(message_3_id.clone(), ["Third message"], cx)
4247        })
4248        .unwrap();
4249    cx.run_until_parked();
4250
4251    // Third message should have second message's input tokens (250) before it
4252    thread.read_with(cx, |thread, _| {
4253        assert_eq!(
4254            thread.tokens_before_message(&message_3_id),
4255            Some(250),
4256            "Third message should have 250 tokens before it (from second request)"
4257        );
4258        // Second message should still have 100
4259        assert_eq!(
4260            thread.tokens_before_message(&message_2_id),
4261            Some(100),
4262            "Second message should still have 100 tokens before it"
4263        );
4264        // First message still has none
4265        assert_eq!(
4266            thread.tokens_before_message(&message_1_id),
4267            None,
4268            "First message should still have no tokens before it"
4269        );
4270    });
4271}
4272
4273#[gpui::test]
4274async fn test_tokens_before_message_after_truncate(cx: &mut TestAppContext) {
4275    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
4276    let fake_model = model.as_fake();
4277
4278    // Set up three messages with responses
4279    let message_1_id = UserMessageId::new();
4280    thread
4281        .update(cx, |thread, cx| {
4282            thread.send(message_1_id.clone(), ["Message 1"], cx)
4283        })
4284        .unwrap();
4285    cx.run_until_parked();
4286    fake_model.send_last_completion_stream_text_chunk("Response 1");
4287    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
4288        language_model::TokenUsage {
4289            input_tokens: 100,
4290            output_tokens: 50,
4291            cache_creation_input_tokens: 0,
4292            cache_read_input_tokens: 0,
4293        },
4294    ));
4295    fake_model.end_last_completion_stream();
4296    cx.run_until_parked();
4297
4298    let message_2_id = UserMessageId::new();
4299    thread
4300        .update(cx, |thread, cx| {
4301            thread.send(message_2_id.clone(), ["Message 2"], cx)
4302        })
4303        .unwrap();
4304    cx.run_until_parked();
4305    fake_model.send_last_completion_stream_text_chunk("Response 2");
4306    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
4307        language_model::TokenUsage {
4308            input_tokens: 250,
4309            output_tokens: 75,
4310            cache_creation_input_tokens: 0,
4311            cache_read_input_tokens: 0,
4312        },
4313    ));
4314    fake_model.end_last_completion_stream();
4315    cx.run_until_parked();
4316
4317    // Verify initial state
4318    thread.read_with(cx, |thread, _| {
4319        assert_eq!(thread.tokens_before_message(&message_2_id), Some(100));
4320    });
4321
4322    // Truncate at message 2 (removes message 2 and everything after)
4323    thread
4324        .update(cx, |thread, cx| thread.truncate(message_2_id.clone(), cx))
4325        .unwrap();
4326    cx.run_until_parked();
4327
4328    // After truncation, message_2_id no longer exists, so lookup should return None
4329    thread.read_with(cx, |thread, _| {
4330        assert_eq!(
4331            thread.tokens_before_message(&message_2_id),
4332            None,
4333            "After truncation, message 2 no longer exists"
4334        );
4335        // Message 1 still exists but has no tokens before it
4336        assert_eq!(
4337            thread.tokens_before_message(&message_1_id),
4338            None,
4339            "First message still has no tokens before it"
4340        );
4341    });
4342}
4343
4344#[gpui::test]
4345async fn test_terminal_tool_permission_rules(cx: &mut TestAppContext) {
4346    init_test(cx);
4347
4348    let fs = FakeFs::new(cx.executor());
4349    fs.insert_tree("/root", json!({})).await;
4350    let project = Project::test(fs, ["/root".as_ref()], cx).await;
4351
4352    // Test 1: Deny rule blocks command
4353    {
4354        let environment = Rc::new(cx.update(|cx| {
4355            FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
4356        }));
4357
4358        cx.update(|cx| {
4359            let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
4360            settings.tool_permissions.tools.insert(
4361                TerminalTool::NAME.into(),
4362                agent_settings::ToolRules {
4363                    default: Some(settings::ToolPermissionMode::Confirm),
4364                    always_allow: vec![],
4365                    always_deny: vec![
4366                        agent_settings::CompiledRegex::new(r"rm\s+-rf", false).unwrap(),
4367                    ],
4368                    always_confirm: vec![],
4369                    invalid_patterns: vec![],
4370                },
4371            );
4372            agent_settings::AgentSettings::override_global(settings, cx);
4373        });
4374
4375        #[allow(clippy::arc_with_non_send_sync)]
4376        let tool = Arc::new(crate::TerminalTool::new(project.clone(), environment));
4377        let (event_stream, _rx) = crate::ToolCallEventStream::test();
4378
4379        let task = cx.update(|cx| {
4380            tool.run(
4381                ToolInput::resolved(crate::TerminalToolInput {
4382                    command: "rm -rf /".to_string(),
4383                    cd: ".".to_string(),
4384                    timeout_ms: None,
4385                }),
4386                event_stream,
4387                cx,
4388            )
4389        });
4390
4391        let result = task.await;
4392        assert!(
4393            result.is_err(),
4394            "expected command to be blocked by deny rule"
4395        );
4396        let err_msg = result.unwrap_err().to_lowercase();
4397        assert!(
4398            err_msg.contains("blocked"),
4399            "error should mention the command was blocked"
4400        );
4401    }
4402
4403    // Test 2: Allow rule skips confirmation (and overrides default: Deny)
4404    {
4405        let environment = Rc::new(cx.update(|cx| {
4406            FakeThreadEnvironment::default()
4407                .with_terminal(FakeTerminalHandle::new_with_immediate_exit(cx, 0))
4408        }));
4409
4410        cx.update(|cx| {
4411            let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
4412            settings.tool_permissions.tools.insert(
4413                TerminalTool::NAME.into(),
4414                agent_settings::ToolRules {
4415                    default: Some(settings::ToolPermissionMode::Deny),
4416                    always_allow: vec![
4417                        agent_settings::CompiledRegex::new(r"^echo\s", false).unwrap(),
4418                    ],
4419                    always_deny: vec![],
4420                    always_confirm: vec![],
4421                    invalid_patterns: vec![],
4422                },
4423            );
4424            agent_settings::AgentSettings::override_global(settings, cx);
4425        });
4426
4427        #[allow(clippy::arc_with_non_send_sync)]
4428        let tool = Arc::new(crate::TerminalTool::new(project.clone(), environment));
4429        let (event_stream, mut rx) = crate::ToolCallEventStream::test();
4430
4431        let task = cx.update(|cx| {
4432            tool.run(
4433                ToolInput::resolved(crate::TerminalToolInput {
4434                    command: "echo hello".to_string(),
4435                    cd: ".".to_string(),
4436                    timeout_ms: None,
4437                }),
4438                event_stream,
4439                cx,
4440            )
4441        });
4442
4443        let update = rx.expect_update_fields().await;
4444        assert!(
4445            update.content.iter().any(|blocks| {
4446                blocks
4447                    .iter()
4448                    .any(|c| matches!(c, acp::ToolCallContent::Terminal(_)))
4449            }),
4450            "expected terminal content (allow rule should skip confirmation and override default deny)"
4451        );
4452
4453        let result = task.await;
4454        assert!(
4455            result.is_ok(),
4456            "expected command to succeed without confirmation"
4457        );
4458    }
4459
4460    // Test 3: global default: allow does NOT override always_confirm patterns
4461    {
4462        let environment = Rc::new(cx.update(|cx| {
4463            FakeThreadEnvironment::default()
4464                .with_terminal(FakeTerminalHandle::new_with_immediate_exit(cx, 0))
4465        }));
4466
4467        cx.update(|cx| {
4468            let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
4469            settings.tool_permissions.default = settings::ToolPermissionMode::Allow;
4470            settings.tool_permissions.tools.insert(
4471                TerminalTool::NAME.into(),
4472                agent_settings::ToolRules {
4473                    default: Some(settings::ToolPermissionMode::Allow),
4474                    always_allow: vec![],
4475                    always_deny: vec![],
4476                    always_confirm: vec![
4477                        agent_settings::CompiledRegex::new(r"sudo", false).unwrap(),
4478                    ],
4479                    invalid_patterns: vec![],
4480                },
4481            );
4482            agent_settings::AgentSettings::override_global(settings, cx);
4483        });
4484
4485        #[allow(clippy::arc_with_non_send_sync)]
4486        let tool = Arc::new(crate::TerminalTool::new(project.clone(), environment));
4487        let (event_stream, mut rx) = crate::ToolCallEventStream::test();
4488
4489        let _task = cx.update(|cx| {
4490            tool.run(
4491                ToolInput::resolved(crate::TerminalToolInput {
4492                    command: "sudo rm file".to_string(),
4493                    cd: ".".to_string(),
4494                    timeout_ms: None,
4495                }),
4496                event_stream,
4497                cx,
4498            )
4499        });
4500
4501        // With global default: allow, confirm patterns are still respected
4502        // The expect_authorization() call will panic if no authorization is requested,
4503        // which validates that the confirm pattern still triggers confirmation
4504        let _auth = rx.expect_authorization().await;
4505
4506        drop(_task);
4507    }
4508
4509    // Test 4: tool-specific default: deny is respected even with global default: allow
4510    {
4511        let environment = Rc::new(cx.update(|cx| {
4512            FakeThreadEnvironment::default()
4513                .with_terminal(FakeTerminalHandle::new_with_immediate_exit(cx, 0))
4514        }));
4515
4516        cx.update(|cx| {
4517            let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
4518            settings.tool_permissions.default = settings::ToolPermissionMode::Allow;
4519            settings.tool_permissions.tools.insert(
4520                TerminalTool::NAME.into(),
4521                agent_settings::ToolRules {
4522                    default: Some(settings::ToolPermissionMode::Deny),
4523                    always_allow: vec![],
4524                    always_deny: vec![],
4525                    always_confirm: vec![],
4526                    invalid_patterns: vec![],
4527                },
4528            );
4529            agent_settings::AgentSettings::override_global(settings, cx);
4530        });
4531
4532        #[allow(clippy::arc_with_non_send_sync)]
4533        let tool = Arc::new(crate::TerminalTool::new(project.clone(), environment));
4534        let (event_stream, _rx) = crate::ToolCallEventStream::test();
4535
4536        let task = cx.update(|cx| {
4537            tool.run(
4538                ToolInput::resolved(crate::TerminalToolInput {
4539                    command: "echo hello".to_string(),
4540                    cd: ".".to_string(),
4541                    timeout_ms: None,
4542                }),
4543                event_stream,
4544                cx,
4545            )
4546        });
4547
4548        // tool-specific default: deny is respected even with global default: allow
4549        let result = task.await;
4550        assert!(
4551            result.is_err(),
4552            "expected command to be blocked by tool-specific deny default"
4553        );
4554        let err_msg = result.unwrap_err().to_lowercase();
4555        assert!(
4556            err_msg.contains("disabled"),
4557            "error should mention the tool is disabled, got: {err_msg}"
4558        );
4559    }
4560}
4561
4562#[gpui::test]
4563async fn test_subagent_tool_call_end_to_end(cx: &mut TestAppContext) {
4564    init_test(cx);
4565    cx.update(|cx| {
4566        LanguageModelRegistry::test(cx);
4567    });
4568    let fs = FakeFs::new(cx.executor());
4569    fs.insert_tree(
4570        "/",
4571        json!({
4572            "a": {
4573                "b.md": "Lorem"
4574            }
4575        }),
4576    )
4577    .await;
4578    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
4579    let thread_store = cx.new(|cx| ThreadStore::new(cx));
4580    let agent = cx.update(|cx| {
4581        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
4582    });
4583    let connection = Rc::new(NativeAgentConnection(agent.clone()));
4584
4585    let acp_thread = cx
4586        .update(|cx| {
4587            connection
4588                .clone()
4589                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
4590        })
4591        .await
4592        .unwrap();
4593    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
4594    let thread = agent.read_with(cx, |agent, _| {
4595        agent.sessions.get(&session_id).unwrap().thread.clone()
4596    });
4597    let model = Arc::new(FakeLanguageModel::default());
4598
4599    // Ensure empty threads are not saved, even if they get mutated.
4600    thread.update(cx, |thread, cx| {
4601        thread.set_model(model.clone(), cx);
4602    });
4603    cx.run_until_parked();
4604
4605    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
4606    cx.run_until_parked();
4607    model.send_last_completion_stream_text_chunk("spawning subagent");
4608    let subagent_tool_input = SpawnAgentToolInput {
4609        label: "label".to_string(),
4610        message: "subagent task prompt".to_string(),
4611        session_id: None,
4612    };
4613    let subagent_tool_use = LanguageModelToolUse {
4614        id: "subagent_1".into(),
4615        name: SpawnAgentTool::NAME.into(),
4616        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
4617        input: serde_json::to_value(&subagent_tool_input).unwrap(),
4618        is_input_complete: true,
4619        thought_signature: None,
4620    };
4621    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
4622        subagent_tool_use,
4623    ));
4624    model.end_last_completion_stream();
4625
4626    cx.run_until_parked();
4627
4628    let subagent_session_id = thread.read_with(cx, |thread, cx| {
4629        thread
4630            .running_subagent_ids(cx)
4631            .get(0)
4632            .expect("subagent thread should be running")
4633            .clone()
4634    });
4635
4636    let subagent_thread = agent.read_with(cx, |agent, _cx| {
4637        agent
4638            .sessions
4639            .get(&subagent_session_id)
4640            .expect("subagent session should exist")
4641            .acp_thread
4642            .clone()
4643    });
4644
4645    model.send_last_completion_stream_text_chunk("subagent task response");
4646    model.end_last_completion_stream();
4647
4648    cx.run_until_parked();
4649
4650    assert_eq!(
4651        subagent_thread.read_with(cx, |thread, cx| thread.to_markdown(cx)),
4652        indoc! {"
4653            ## User
4654
4655            subagent task prompt
4656
4657            ## Assistant
4658
4659            subagent task response
4660
4661        "}
4662    );
4663
4664    model.send_last_completion_stream_text_chunk("Response");
4665    model.end_last_completion_stream();
4666
4667    send.await.unwrap();
4668
4669    assert_eq!(
4670        acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx)),
4671        indoc! {r#"
4672            ## User
4673
4674            Prompt
4675
4676            ## Assistant
4677
4678            spawning subagent
4679
4680            **Tool Call: label**
4681            Status: Completed
4682
4683            subagent task response
4684
4685            ## Assistant
4686
4687            Response
4688
4689        "#},
4690    );
4691}
4692
4693#[gpui::test]
4694async fn test_subagent_tool_output_does_not_include_thinking(cx: &mut TestAppContext) {
4695    init_test(cx);
4696    cx.update(|cx| {
4697        LanguageModelRegistry::test(cx);
4698    });
4699    let fs = FakeFs::new(cx.executor());
4700    fs.insert_tree(
4701        "/",
4702        json!({
4703            "a": {
4704                "b.md": "Lorem"
4705            }
4706        }),
4707    )
4708    .await;
4709    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
4710    let thread_store = cx.new(|cx| ThreadStore::new(cx));
4711    let agent = cx.update(|cx| {
4712        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
4713    });
4714    let connection = Rc::new(NativeAgentConnection(agent.clone()));
4715
4716    let acp_thread = cx
4717        .update(|cx| {
4718            connection
4719                .clone()
4720                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
4721        })
4722        .await
4723        .unwrap();
4724    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
4725    let thread = agent.read_with(cx, |agent, _| {
4726        agent.sessions.get(&session_id).unwrap().thread.clone()
4727    });
4728    let model = Arc::new(FakeLanguageModel::default());
4729
4730    // Ensure empty threads are not saved, even if they get mutated.
4731    thread.update(cx, |thread, cx| {
4732        thread.set_model(model.clone(), cx);
4733    });
4734    cx.run_until_parked();
4735
4736    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
4737    cx.run_until_parked();
4738    model.send_last_completion_stream_text_chunk("spawning subagent");
4739    let subagent_tool_input = SpawnAgentToolInput {
4740        label: "label".to_string(),
4741        message: "subagent task prompt".to_string(),
4742        session_id: None,
4743    };
4744    let subagent_tool_use = LanguageModelToolUse {
4745        id: "subagent_1".into(),
4746        name: SpawnAgentTool::NAME.into(),
4747        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
4748        input: serde_json::to_value(&subagent_tool_input).unwrap(),
4749        is_input_complete: true,
4750        thought_signature: None,
4751    };
4752    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
4753        subagent_tool_use,
4754    ));
4755    model.end_last_completion_stream();
4756
4757    cx.run_until_parked();
4758
4759    let subagent_session_id = thread.read_with(cx, |thread, cx| {
4760        thread
4761            .running_subagent_ids(cx)
4762            .get(0)
4763            .expect("subagent thread should be running")
4764            .clone()
4765    });
4766
4767    let subagent_thread = agent.read_with(cx, |agent, _cx| {
4768        agent
4769            .sessions
4770            .get(&subagent_session_id)
4771            .expect("subagent session should exist")
4772            .acp_thread
4773            .clone()
4774    });
4775
4776    model.send_last_completion_stream_text_chunk("subagent task response 1");
4777    model.send_last_completion_stream_event(LanguageModelCompletionEvent::Thinking {
4778        text: "thinking more about the subagent task".into(),
4779        signature: None,
4780    });
4781    model.send_last_completion_stream_text_chunk("subagent task response 2");
4782    model.end_last_completion_stream();
4783
4784    cx.run_until_parked();
4785
4786    assert_eq!(
4787        subagent_thread.read_with(cx, |thread, cx| thread.to_markdown(cx)),
4788        indoc! {"
4789            ## User
4790
4791            subagent task prompt
4792
4793            ## Assistant
4794
4795            subagent task response 1
4796
4797            <thinking>
4798            thinking more about the subagent task
4799            </thinking>
4800
4801            subagent task response 2
4802
4803        "}
4804    );
4805
4806    model.send_last_completion_stream_text_chunk("Response");
4807    model.end_last_completion_stream();
4808
4809    send.await.unwrap();
4810
4811    assert_eq!(
4812        acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx)),
4813        indoc! {r#"
4814            ## User
4815
4816            Prompt
4817
4818            ## Assistant
4819
4820            spawning subagent
4821
4822            **Tool Call: label**
4823            Status: Completed
4824
4825            subagent task response 1
4826
4827            subagent task response 2
4828
4829            ## Assistant
4830
4831            Response
4832
4833        "#},
4834    );
4835}
4836
4837#[gpui::test]
4838async fn test_subagent_tool_call_cancellation_during_task_prompt(cx: &mut TestAppContext) {
4839    init_test(cx);
4840    cx.update(|cx| {
4841        LanguageModelRegistry::test(cx);
4842    });
4843    let fs = FakeFs::new(cx.executor());
4844    fs.insert_tree(
4845        "/",
4846        json!({
4847            "a": {
4848                "b.md": "Lorem"
4849            }
4850        }),
4851    )
4852    .await;
4853    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
4854    let thread_store = cx.new(|cx| ThreadStore::new(cx));
4855    let agent = cx.update(|cx| {
4856        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
4857    });
4858    let connection = Rc::new(NativeAgentConnection(agent.clone()));
4859
4860    let acp_thread = cx
4861        .update(|cx| {
4862            connection
4863                .clone()
4864                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
4865        })
4866        .await
4867        .unwrap();
4868    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
4869    let thread = agent.read_with(cx, |agent, _| {
4870        agent.sessions.get(&session_id).unwrap().thread.clone()
4871    });
4872    let model = Arc::new(FakeLanguageModel::default());
4873
4874    // Ensure empty threads are not saved, even if they get mutated.
4875    thread.update(cx, |thread, cx| {
4876        thread.set_model(model.clone(), cx);
4877    });
4878    cx.run_until_parked();
4879
4880    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
4881    cx.run_until_parked();
4882    model.send_last_completion_stream_text_chunk("spawning subagent");
4883    let subagent_tool_input = SpawnAgentToolInput {
4884        label: "label".to_string(),
4885        message: "subagent task prompt".to_string(),
4886        session_id: None,
4887    };
4888    let subagent_tool_use = LanguageModelToolUse {
4889        id: "subagent_1".into(),
4890        name: SpawnAgentTool::NAME.into(),
4891        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
4892        input: serde_json::to_value(&subagent_tool_input).unwrap(),
4893        is_input_complete: true,
4894        thought_signature: None,
4895    };
4896    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
4897        subagent_tool_use,
4898    ));
4899    model.end_last_completion_stream();
4900
4901    cx.run_until_parked();
4902
4903    let subagent_session_id = thread.read_with(cx, |thread, cx| {
4904        thread
4905            .running_subagent_ids(cx)
4906            .get(0)
4907            .expect("subagent thread should be running")
4908            .clone()
4909    });
4910    let subagent_acp_thread = agent.read_with(cx, |agent, _cx| {
4911        agent
4912            .sessions
4913            .get(&subagent_session_id)
4914            .expect("subagent session should exist")
4915            .acp_thread
4916            .clone()
4917    });
4918
4919    // model.send_last_completion_stream_text_chunk("subagent task response");
4920    // model.end_last_completion_stream();
4921
4922    // cx.run_until_parked();
4923
4924    acp_thread.update(cx, |thread, cx| thread.cancel(cx)).await;
4925
4926    cx.run_until_parked();
4927
4928    send.await.unwrap();
4929
4930    acp_thread.read_with(cx, |thread, cx| {
4931        assert_eq!(thread.status(), ThreadStatus::Idle);
4932        assert_eq!(
4933            thread.to_markdown(cx),
4934            indoc! {"
4935                ## User
4936
4937                Prompt
4938
4939                ## Assistant
4940
4941                spawning subagent
4942
4943                **Tool Call: label**
4944                Status: Canceled
4945
4946            "}
4947        );
4948    });
4949    subagent_acp_thread.read_with(cx, |thread, cx| {
4950        assert_eq!(thread.status(), ThreadStatus::Idle);
4951        assert_eq!(
4952            thread.to_markdown(cx),
4953            indoc! {"
4954                ## User
4955
4956                subagent task prompt
4957
4958            "}
4959        );
4960    });
4961}
4962
4963#[gpui::test]
4964async fn test_subagent_tool_resume_session(cx: &mut TestAppContext) {
4965    init_test(cx);
4966    cx.update(|cx| {
4967        LanguageModelRegistry::test(cx);
4968    });
4969    let fs = FakeFs::new(cx.executor());
4970    fs.insert_tree(
4971        "/",
4972        json!({
4973            "a": {
4974                "b.md": "Lorem"
4975            }
4976        }),
4977    )
4978    .await;
4979    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
4980    let thread_store = cx.new(|cx| ThreadStore::new(cx));
4981    let agent = cx.update(|cx| {
4982        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
4983    });
4984    let connection = Rc::new(NativeAgentConnection(agent.clone()));
4985
4986    let acp_thread = cx
4987        .update(|cx| {
4988            connection
4989                .clone()
4990                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
4991        })
4992        .await
4993        .unwrap();
4994    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
4995    let thread = agent.read_with(cx, |agent, _| {
4996        agent.sessions.get(&session_id).unwrap().thread.clone()
4997    });
4998    let model = Arc::new(FakeLanguageModel::default());
4999
5000    thread.update(cx, |thread, cx| {
5001        thread.set_model(model.clone(), cx);
5002    });
5003    cx.run_until_parked();
5004
5005    // === First turn: create subagent ===
5006    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("First prompt", cx));
5007    cx.run_until_parked();
5008    model.send_last_completion_stream_text_chunk("spawning subagent");
5009    let subagent_tool_input = SpawnAgentToolInput {
5010        label: "initial task".to_string(),
5011        message: "do the first task".to_string(),
5012        session_id: None,
5013    };
5014    let subagent_tool_use = LanguageModelToolUse {
5015        id: "subagent_1".into(),
5016        name: SpawnAgentTool::NAME.into(),
5017        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
5018        input: serde_json::to_value(&subagent_tool_input).unwrap(),
5019        is_input_complete: true,
5020        thought_signature: None,
5021    };
5022    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
5023        subagent_tool_use,
5024    ));
5025    model.end_last_completion_stream();
5026
5027    cx.run_until_parked();
5028
5029    let subagent_session_id = thread.read_with(cx, |thread, cx| {
5030        thread
5031            .running_subagent_ids(cx)
5032            .get(0)
5033            .expect("subagent thread should be running")
5034            .clone()
5035    });
5036
5037    let subagent_acp_thread = agent.read_with(cx, |agent, _cx| {
5038        agent
5039            .sessions
5040            .get(&subagent_session_id)
5041            .expect("subagent session should exist")
5042            .acp_thread
5043            .clone()
5044    });
5045
5046    // Subagent responds
5047    model.send_last_completion_stream_text_chunk("first task response");
5048    model.end_last_completion_stream();
5049
5050    cx.run_until_parked();
5051
5052    // Parent model responds to complete first turn
5053    model.send_last_completion_stream_text_chunk("First response");
5054    model.end_last_completion_stream();
5055
5056    send.await.unwrap();
5057
5058    // Verify subagent is no longer running
5059    thread.read_with(cx, |thread, cx| {
5060        assert!(
5061            thread.running_subagent_ids(cx).is_empty(),
5062            "subagent should not be running after completion"
5063        );
5064    });
5065
5066    // === Second turn: resume subagent with session_id ===
5067    let send2 = acp_thread.update(cx, |thread, cx| thread.send_raw("Follow up", cx));
5068    cx.run_until_parked();
5069    model.send_last_completion_stream_text_chunk("resuming subagent");
5070    let resume_tool_input = SpawnAgentToolInput {
5071        label: "follow-up task".to_string(),
5072        message: "do the follow-up task".to_string(),
5073        session_id: Some(subagent_session_id.clone()),
5074    };
5075    let resume_tool_use = LanguageModelToolUse {
5076        id: "subagent_2".into(),
5077        name: SpawnAgentTool::NAME.into(),
5078        raw_input: serde_json::to_string(&resume_tool_input).unwrap(),
5079        input: serde_json::to_value(&resume_tool_input).unwrap(),
5080        is_input_complete: true,
5081        thought_signature: None,
5082    };
5083    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(resume_tool_use));
5084    model.end_last_completion_stream();
5085
5086    cx.run_until_parked();
5087
5088    // Subagent should be running again with the same session
5089    thread.read_with(cx, |thread, cx| {
5090        let running = thread.running_subagent_ids(cx);
5091        assert_eq!(running.len(), 1, "subagent should be running");
5092        assert_eq!(running[0], subagent_session_id, "should be same session");
5093    });
5094
5095    // Subagent responds to follow-up
5096    model.send_last_completion_stream_text_chunk("follow-up task response");
5097    model.end_last_completion_stream();
5098
5099    cx.run_until_parked();
5100
5101    // Parent model responds to complete second turn
5102    model.send_last_completion_stream_text_chunk("Second response");
5103    model.end_last_completion_stream();
5104
5105    send2.await.unwrap();
5106
5107    // Verify subagent is no longer running
5108    thread.read_with(cx, |thread, cx| {
5109        assert!(
5110            thread.running_subagent_ids(cx).is_empty(),
5111            "subagent should not be running after resume completion"
5112        );
5113    });
5114
5115    // Verify the subagent's acp thread has both conversation turns
5116    assert_eq!(
5117        subagent_acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx)),
5118        indoc! {"
5119            ## User
5120
5121            do the first task
5122
5123            ## Assistant
5124
5125            first task response
5126
5127            ## User
5128
5129            do the follow-up task
5130
5131            ## Assistant
5132
5133            follow-up task response
5134
5135        "}
5136    );
5137}
5138
5139#[gpui::test]
5140async fn test_subagent_tool_is_present_when_feature_flag_enabled(cx: &mut TestAppContext) {
5141    init_test(cx);
5142
5143    let fs = FakeFs::new(cx.executor());
5144    fs.insert_tree(path!("/test"), json!({})).await;
5145    let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
5146    let project_context = cx.new(|_cx| ProjectContext::default());
5147    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
5148    let context_server_registry =
5149        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
5150    let model = Arc::new(FakeLanguageModel::default());
5151
5152    let environment = Rc::new(cx.update(|cx| {
5153        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
5154    }));
5155
5156    let thread = cx.new(|cx| {
5157        let mut thread = Thread::new(
5158            project.clone(),
5159            project_context,
5160            context_server_registry,
5161            Templates::new(),
5162            Some(model),
5163            cx,
5164        );
5165        thread.add_default_tools(environment, cx);
5166        thread
5167    });
5168
5169    thread.read_with(cx, |thread, _| {
5170        assert!(
5171            thread.has_registered_tool(SpawnAgentTool::NAME),
5172            "subagent tool should be present when feature flag is enabled"
5173        );
5174    });
5175}
5176
5177#[gpui::test]
5178async fn test_subagent_thread_inherits_parent_thread_properties(cx: &mut TestAppContext) {
5179    init_test(cx);
5180
5181    let fs = FakeFs::new(cx.executor());
5182    fs.insert_tree(path!("/test"), json!({})).await;
5183    let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
5184    let project_context = cx.new(|_cx| ProjectContext::default());
5185    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
5186    let context_server_registry =
5187        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
5188    let model = Arc::new(FakeLanguageModel::default());
5189
5190    let parent_thread = cx.new(|cx| {
5191        Thread::new(
5192            project.clone(),
5193            project_context,
5194            context_server_registry,
5195            Templates::new(),
5196            Some(model.clone()),
5197            cx,
5198        )
5199    });
5200
5201    let subagent_thread = cx.new(|cx| Thread::new_subagent(&parent_thread, cx));
5202    subagent_thread.read_with(cx, |subagent_thread, cx| {
5203        assert!(subagent_thread.is_subagent());
5204        assert_eq!(subagent_thread.depth(), 1);
5205        assert_eq!(
5206            subagent_thread.model().map(|model| model.id()),
5207            Some(model.id())
5208        );
5209        assert_eq!(
5210            subagent_thread.parent_thread_id(),
5211            Some(parent_thread.read(cx).id().clone())
5212        );
5213
5214        let request = subagent_thread
5215            .build_completion_request(CompletionIntent::UserPrompt, cx)
5216            .unwrap();
5217        assert_eq!(request.intent, Some(CompletionIntent::Subagent));
5218    });
5219}
5220
5221#[gpui::test]
5222async fn test_max_subagent_depth_prevents_tool_registration(cx: &mut TestAppContext) {
5223    init_test(cx);
5224
5225    let fs = FakeFs::new(cx.executor());
5226    fs.insert_tree(path!("/test"), json!({})).await;
5227    let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
5228    let project_context = cx.new(|_cx| ProjectContext::default());
5229    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
5230    let context_server_registry =
5231        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
5232    let model = Arc::new(FakeLanguageModel::default());
5233    let environment = Rc::new(cx.update(|cx| {
5234        FakeThreadEnvironment::default().with_terminal(FakeTerminalHandle::new_never_exits(cx))
5235    }));
5236
5237    let deep_parent_thread = cx.new(|cx| {
5238        let mut thread = Thread::new(
5239            project.clone(),
5240            project_context,
5241            context_server_registry,
5242            Templates::new(),
5243            Some(model.clone()),
5244            cx,
5245        );
5246        thread.set_subagent_context(SubagentContext {
5247            parent_thread_id: agent_client_protocol::SessionId::new("parent-id"),
5248            depth: MAX_SUBAGENT_DEPTH - 1,
5249        });
5250        thread
5251    });
5252    let deep_subagent_thread = cx.new(|cx| {
5253        let mut thread = Thread::new_subagent(&deep_parent_thread, cx);
5254        thread.add_default_tools(environment, cx);
5255        thread
5256    });
5257
5258    deep_subagent_thread.read_with(cx, |thread, _| {
5259        assert_eq!(thread.depth(), MAX_SUBAGENT_DEPTH);
5260        assert!(
5261            !thread.has_registered_tool(SpawnAgentTool::NAME),
5262            "subagent tool should not be present at max depth"
5263        );
5264    });
5265}
5266
5267#[gpui::test]
5268async fn test_parent_cancel_stops_subagent(cx: &mut TestAppContext) {
5269    init_test(cx);
5270
5271    let fs = FakeFs::new(cx.executor());
5272    fs.insert_tree(path!("/test"), json!({})).await;
5273    let project = Project::test(fs, [path!("/test").as_ref()], cx).await;
5274    let project_context = cx.new(|_cx| ProjectContext::default());
5275    let context_server_store = project.read_with(cx, |project, _| project.context_server_store());
5276    let context_server_registry =
5277        cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
5278    let model = Arc::new(FakeLanguageModel::default());
5279
5280    let parent = cx.new(|cx| {
5281        Thread::new(
5282            project.clone(),
5283            project_context.clone(),
5284            context_server_registry.clone(),
5285            Templates::new(),
5286            Some(model.clone()),
5287            cx,
5288        )
5289    });
5290
5291    let subagent = cx.new(|cx| Thread::new_subagent(&parent, cx));
5292
5293    parent.update(cx, |thread, _cx| {
5294        thread.register_running_subagent(subagent.downgrade());
5295    });
5296
5297    subagent
5298        .update(cx, |thread, cx| {
5299            thread.send(UserMessageId::new(), ["Do work".to_string()], cx)
5300        })
5301        .unwrap();
5302    cx.run_until_parked();
5303
5304    subagent.read_with(cx, |thread, _| {
5305        assert!(!thread.is_turn_complete(), "subagent should be running");
5306    });
5307
5308    parent.update(cx, |thread, cx| {
5309        thread.cancel(cx).detach();
5310    });
5311
5312    subagent.read_with(cx, |thread, _| {
5313        assert!(
5314            thread.is_turn_complete(),
5315            "subagent should be cancelled when parent cancels"
5316        );
5317    });
5318}
5319
5320#[gpui::test]
5321async fn test_subagent_context_window_warning(cx: &mut TestAppContext) {
5322    init_test(cx);
5323    cx.update(|cx| {
5324        LanguageModelRegistry::test(cx);
5325    });
5326    let fs = FakeFs::new(cx.executor());
5327    fs.insert_tree(
5328        "/",
5329        json!({
5330            "a": {
5331                "b.md": "Lorem"
5332            }
5333        }),
5334    )
5335    .await;
5336    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
5337    let thread_store = cx.new(|cx| ThreadStore::new(cx));
5338    let agent = cx.update(|cx| {
5339        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
5340    });
5341    let connection = Rc::new(NativeAgentConnection(agent.clone()));
5342
5343    let acp_thread = cx
5344        .update(|cx| {
5345            connection
5346                .clone()
5347                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
5348        })
5349        .await
5350        .unwrap();
5351    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
5352    let thread = agent.read_with(cx, |agent, _| {
5353        agent.sessions.get(&session_id).unwrap().thread.clone()
5354    });
5355    let model = Arc::new(FakeLanguageModel::default());
5356
5357    thread.update(cx, |thread, cx| {
5358        thread.set_model(model.clone(), cx);
5359    });
5360    cx.run_until_parked();
5361
5362    // Start the parent turn
5363    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
5364    cx.run_until_parked();
5365    model.send_last_completion_stream_text_chunk("spawning subagent");
5366    let subagent_tool_input = SpawnAgentToolInput {
5367        label: "label".to_string(),
5368        message: "subagent task prompt".to_string(),
5369        session_id: None,
5370    };
5371    let subagent_tool_use = LanguageModelToolUse {
5372        id: "subagent_1".into(),
5373        name: SpawnAgentTool::NAME.into(),
5374        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
5375        input: serde_json::to_value(&subagent_tool_input).unwrap(),
5376        is_input_complete: true,
5377        thought_signature: None,
5378    };
5379    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
5380        subagent_tool_use,
5381    ));
5382    model.end_last_completion_stream();
5383
5384    cx.run_until_parked();
5385
5386    // Verify subagent is running
5387    let subagent_session_id = thread.read_with(cx, |thread, cx| {
5388        thread
5389            .running_subagent_ids(cx)
5390            .get(0)
5391            .expect("subagent thread should be running")
5392            .clone()
5393    });
5394
5395    // Send a usage update that crosses the warning threshold (80% of 1,000,000)
5396    model.send_last_completion_stream_text_chunk("partial work");
5397    model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
5398        TokenUsage {
5399            input_tokens: 850_000,
5400            output_tokens: 0,
5401            cache_creation_input_tokens: 0,
5402            cache_read_input_tokens: 0,
5403        },
5404    ));
5405
5406    cx.run_until_parked();
5407
5408    // The subagent should no longer be running
5409    thread.read_with(cx, |thread, cx| {
5410        assert!(
5411            thread.running_subagent_ids(cx).is_empty(),
5412            "subagent should be stopped after context window warning"
5413        );
5414    });
5415
5416    // The parent model should get a new completion request to respond to the tool error
5417    model.send_last_completion_stream_text_chunk("Response after warning");
5418    model.end_last_completion_stream();
5419
5420    send.await.unwrap();
5421
5422    // Verify the parent thread shows the warning error in the tool call
5423    let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
5424    assert!(
5425        markdown.contains("nearing the end of its context window"),
5426        "tool output should contain context window warning message, got:\n{markdown}"
5427    );
5428    assert!(
5429        markdown.contains("Status: Failed"),
5430        "tool call should have Failed status, got:\n{markdown}"
5431    );
5432
5433    // Verify the subagent session still exists (can be resumed)
5434    agent.read_with(cx, |agent, _cx| {
5435        assert!(
5436            agent.sessions.contains_key(&subagent_session_id),
5437            "subagent session should still exist for potential resume"
5438        );
5439    });
5440}
5441
5442#[gpui::test]
5443async fn test_subagent_no_context_window_warning_when_already_at_warning(cx: &mut TestAppContext) {
5444    init_test(cx);
5445    cx.update(|cx| {
5446        LanguageModelRegistry::test(cx);
5447    });
5448    let fs = FakeFs::new(cx.executor());
5449    fs.insert_tree(
5450        "/",
5451        json!({
5452            "a": {
5453                "b.md": "Lorem"
5454            }
5455        }),
5456    )
5457    .await;
5458    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
5459    let thread_store = cx.new(|cx| ThreadStore::new(cx));
5460    let agent = cx.update(|cx| {
5461        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
5462    });
5463    let connection = Rc::new(NativeAgentConnection(agent.clone()));
5464
5465    let acp_thread = cx
5466        .update(|cx| {
5467            connection
5468                .clone()
5469                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
5470        })
5471        .await
5472        .unwrap();
5473    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
5474    let thread = agent.read_with(cx, |agent, _| {
5475        agent.sessions.get(&session_id).unwrap().thread.clone()
5476    });
5477    let model = Arc::new(FakeLanguageModel::default());
5478
5479    thread.update(cx, |thread, cx| {
5480        thread.set_model(model.clone(), cx);
5481    });
5482    cx.run_until_parked();
5483
5484    // === First turn: create subagent, trigger context window warning ===
5485    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("First prompt", cx));
5486    cx.run_until_parked();
5487    model.send_last_completion_stream_text_chunk("spawning subagent");
5488    let subagent_tool_input = SpawnAgentToolInput {
5489        label: "initial task".to_string(),
5490        message: "do the first task".to_string(),
5491        session_id: None,
5492    };
5493    let subagent_tool_use = LanguageModelToolUse {
5494        id: "subagent_1".into(),
5495        name: SpawnAgentTool::NAME.into(),
5496        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
5497        input: serde_json::to_value(&subagent_tool_input).unwrap(),
5498        is_input_complete: true,
5499        thought_signature: None,
5500    };
5501    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
5502        subagent_tool_use,
5503    ));
5504    model.end_last_completion_stream();
5505
5506    cx.run_until_parked();
5507
5508    let subagent_session_id = thread.read_with(cx, |thread, cx| {
5509        thread
5510            .running_subagent_ids(cx)
5511            .get(0)
5512            .expect("subagent thread should be running")
5513            .clone()
5514    });
5515
5516    // Subagent sends a usage update that crosses the warning threshold.
5517    // This triggers Normal→Warning, stopping the subagent.
5518    model.send_last_completion_stream_text_chunk("partial work");
5519    model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
5520        TokenUsage {
5521            input_tokens: 850_000,
5522            output_tokens: 0,
5523            cache_creation_input_tokens: 0,
5524            cache_read_input_tokens: 0,
5525        },
5526    ));
5527
5528    cx.run_until_parked();
5529
5530    // Verify the first turn was stopped with a context window warning
5531    thread.read_with(cx, |thread, cx| {
5532        assert!(
5533            thread.running_subagent_ids(cx).is_empty(),
5534            "subagent should be stopped after context window warning"
5535        );
5536    });
5537
5538    // Parent model responds to complete first turn
5539    model.send_last_completion_stream_text_chunk("First response");
5540    model.end_last_completion_stream();
5541
5542    send.await.unwrap();
5543
5544    let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
5545    assert!(
5546        markdown.contains("nearing the end of its context window"),
5547        "first turn should have context window warning, got:\n{markdown}"
5548    );
5549
5550    // === Second turn: resume the same subagent (now at Warning level) ===
5551    let send2 = acp_thread.update(cx, |thread, cx| thread.send_raw("Follow up", cx));
5552    cx.run_until_parked();
5553    model.send_last_completion_stream_text_chunk("resuming subagent");
5554    let resume_tool_input = SpawnAgentToolInput {
5555        label: "follow-up task".to_string(),
5556        message: "do the follow-up task".to_string(),
5557        session_id: Some(subagent_session_id.clone()),
5558    };
5559    let resume_tool_use = LanguageModelToolUse {
5560        id: "subagent_2".into(),
5561        name: SpawnAgentTool::NAME.into(),
5562        raw_input: serde_json::to_string(&resume_tool_input).unwrap(),
5563        input: serde_json::to_value(&resume_tool_input).unwrap(),
5564        is_input_complete: true,
5565        thought_signature: None,
5566    };
5567    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(resume_tool_use));
5568    model.end_last_completion_stream();
5569
5570    cx.run_until_parked();
5571
5572    // Subagent responds with tokens still at warning level (no worse).
5573    // Since ratio_before_prompt was already Warning, this should NOT
5574    // trigger the context window warning again.
5575    model.send_last_completion_stream_text_chunk("follow-up task response");
5576    model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
5577        TokenUsage {
5578            input_tokens: 870_000,
5579            output_tokens: 0,
5580            cache_creation_input_tokens: 0,
5581            cache_read_input_tokens: 0,
5582        },
5583    ));
5584    model.end_last_completion_stream();
5585
5586    cx.run_until_parked();
5587
5588    // Parent model responds to complete second turn
5589    model.send_last_completion_stream_text_chunk("Second response");
5590    model.end_last_completion_stream();
5591
5592    send2.await.unwrap();
5593
5594    // The resumed subagent should have completed normally since the ratio
5595    // didn't transition (it was Warning before and stayed at Warning)
5596    let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
5597    assert!(
5598        markdown.contains("follow-up task response"),
5599        "resumed subagent should complete normally when already at warning, got:\n{markdown}"
5600    );
5601    // The second tool call should NOT have a context window warning
5602    let second_tool_pos = markdown
5603        .find("follow-up task")
5604        .expect("should find follow-up tool call");
5605    let after_second_tool = &markdown[second_tool_pos..];
5606    assert!(
5607        !after_second_tool.contains("nearing the end of its context window"),
5608        "should NOT contain context window warning for resumed subagent at same level, got:\n{after_second_tool}"
5609    );
5610}
5611
5612#[gpui::test]
5613async fn test_subagent_error_propagation(cx: &mut TestAppContext) {
5614    init_test(cx);
5615    cx.update(|cx| {
5616        LanguageModelRegistry::test(cx);
5617    });
5618    let fs = FakeFs::new(cx.executor());
5619    fs.insert_tree(
5620        "/",
5621        json!({
5622            "a": {
5623                "b.md": "Lorem"
5624            }
5625        }),
5626    )
5627    .await;
5628    let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
5629    let thread_store = cx.new(|cx| ThreadStore::new(cx));
5630    let agent = cx.update(|cx| {
5631        NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
5632    });
5633    let connection = Rc::new(NativeAgentConnection(agent.clone()));
5634
5635    let acp_thread = cx
5636        .update(|cx| {
5637            connection
5638                .clone()
5639                .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
5640        })
5641        .await
5642        .unwrap();
5643    let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
5644    let thread = agent.read_with(cx, |agent, _| {
5645        agent.sessions.get(&session_id).unwrap().thread.clone()
5646    });
5647    let model = Arc::new(FakeLanguageModel::default());
5648
5649    thread.update(cx, |thread, cx| {
5650        thread.set_model(model.clone(), cx);
5651    });
5652    cx.run_until_parked();
5653
5654    // Start the parent turn
5655    let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
5656    cx.run_until_parked();
5657    model.send_last_completion_stream_text_chunk("spawning subagent");
5658    let subagent_tool_input = SpawnAgentToolInput {
5659        label: "label".to_string(),
5660        message: "subagent task prompt".to_string(),
5661        session_id: None,
5662    };
5663    let subagent_tool_use = LanguageModelToolUse {
5664        id: "subagent_1".into(),
5665        name: SpawnAgentTool::NAME.into(),
5666        raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
5667        input: serde_json::to_value(&subagent_tool_input).unwrap(),
5668        is_input_complete: true,
5669        thought_signature: None,
5670    };
5671    model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
5672        subagent_tool_use,
5673    ));
5674    model.end_last_completion_stream();
5675
5676    cx.run_until_parked();
5677
5678    // Verify subagent is running
5679    thread.read_with(cx, |thread, cx| {
5680        assert!(
5681            !thread.running_subagent_ids(cx).is_empty(),
5682            "subagent should be running"
5683        );
5684    });
5685
5686    // The subagent's model returns a non-retryable error
5687    model.send_last_completion_stream_error(LanguageModelCompletionError::PromptTooLarge {
5688        tokens: None,
5689    });
5690
5691    cx.run_until_parked();
5692
5693    // The subagent should no longer be running
5694    thread.read_with(cx, |thread, cx| {
5695        assert!(
5696            thread.running_subagent_ids(cx).is_empty(),
5697            "subagent should not be running after error"
5698        );
5699    });
5700
5701    // The parent model should get a new completion request to respond to the tool error
5702    model.send_last_completion_stream_text_chunk("Response after error");
5703    model.end_last_completion_stream();
5704
5705    send.await.unwrap();
5706
5707    // Verify the parent thread shows the error in the tool call
5708    let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
5709    assert!(
5710        markdown.contains("Status: Failed"),
5711        "tool call should have Failed status after model error, got:\n{markdown}"
5712    );
5713}
5714
5715#[gpui::test]
5716async fn test_edit_file_tool_deny_rule_blocks_edit(cx: &mut TestAppContext) {
5717    init_test(cx);
5718
5719    let fs = FakeFs::new(cx.executor());
5720    fs.insert_tree("/root", json!({"sensitive_config.txt": "secret data"}))
5721        .await;
5722    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
5723
5724    cx.update(|cx| {
5725        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
5726        settings.tool_permissions.tools.insert(
5727            EditFileTool::NAME.into(),
5728            agent_settings::ToolRules {
5729                default: Some(settings::ToolPermissionMode::Allow),
5730                always_allow: vec![],
5731                always_deny: vec![agent_settings::CompiledRegex::new(r"sensitive", false).unwrap()],
5732                always_confirm: vec![],
5733                invalid_patterns: vec![],
5734            },
5735        );
5736        agent_settings::AgentSettings::override_global(settings, cx);
5737    });
5738
5739    let context_server_registry =
5740        cx.new(|cx| crate::ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
5741    let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
5742    let templates = crate::Templates::new();
5743    let thread = cx.new(|cx| {
5744        crate::Thread::new(
5745            project.clone(),
5746            cx.new(|_cx| prompt_store::ProjectContext::default()),
5747            context_server_registry,
5748            templates.clone(),
5749            None,
5750            cx,
5751        )
5752    });
5753
5754    #[allow(clippy::arc_with_non_send_sync)]
5755    let tool = Arc::new(crate::EditFileTool::new(
5756        project.clone(),
5757        thread.downgrade(),
5758        language_registry,
5759        templates,
5760    ));
5761    let (event_stream, _rx) = crate::ToolCallEventStream::test();
5762
5763    let task = cx.update(|cx| {
5764        tool.run(
5765            ToolInput::resolved(crate::EditFileToolInput {
5766                display_description: "Edit sensitive file".to_string(),
5767                path: "root/sensitive_config.txt".into(),
5768                mode: crate::EditFileMode::Edit,
5769            }),
5770            event_stream,
5771            cx,
5772        )
5773    });
5774
5775    let result = task.await;
5776    assert!(result.is_err(), "expected edit to be blocked");
5777    assert!(
5778        result.unwrap_err().to_string().contains("blocked"),
5779        "error should mention the edit was blocked"
5780    );
5781}
5782
5783#[gpui::test]
5784async fn test_delete_path_tool_deny_rule_blocks_deletion(cx: &mut TestAppContext) {
5785    init_test(cx);
5786
5787    let fs = FakeFs::new(cx.executor());
5788    fs.insert_tree("/root", json!({"important_data.txt": "critical info"}))
5789        .await;
5790    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
5791
5792    cx.update(|cx| {
5793        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
5794        settings.tool_permissions.tools.insert(
5795            DeletePathTool::NAME.into(),
5796            agent_settings::ToolRules {
5797                default: Some(settings::ToolPermissionMode::Allow),
5798                always_allow: vec![],
5799                always_deny: vec![agent_settings::CompiledRegex::new(r"important", false).unwrap()],
5800                always_confirm: vec![],
5801                invalid_patterns: vec![],
5802            },
5803        );
5804        agent_settings::AgentSettings::override_global(settings, cx);
5805    });
5806
5807    let action_log = cx.new(|_cx| action_log::ActionLog::new(project.clone()));
5808
5809    #[allow(clippy::arc_with_non_send_sync)]
5810    let tool = Arc::new(crate::DeletePathTool::new(project, action_log));
5811    let (event_stream, _rx) = crate::ToolCallEventStream::test();
5812
5813    let task = cx.update(|cx| {
5814        tool.run(
5815            ToolInput::resolved(crate::DeletePathToolInput {
5816                path: "root/important_data.txt".to_string(),
5817            }),
5818            event_stream,
5819            cx,
5820        )
5821    });
5822
5823    let result = task.await;
5824    assert!(result.is_err(), "expected deletion to be blocked");
5825    assert!(
5826        result.unwrap_err().contains("blocked"),
5827        "error should mention the deletion was blocked"
5828    );
5829}
5830
5831#[gpui::test]
5832async fn test_move_path_tool_denies_if_destination_denied(cx: &mut TestAppContext) {
5833    init_test(cx);
5834
5835    let fs = FakeFs::new(cx.executor());
5836    fs.insert_tree(
5837        "/root",
5838        json!({
5839            "safe.txt": "content",
5840            "protected": {}
5841        }),
5842    )
5843    .await;
5844    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
5845
5846    cx.update(|cx| {
5847        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
5848        settings.tool_permissions.tools.insert(
5849            MovePathTool::NAME.into(),
5850            agent_settings::ToolRules {
5851                default: Some(settings::ToolPermissionMode::Allow),
5852                always_allow: vec![],
5853                always_deny: vec![agent_settings::CompiledRegex::new(r"protected", false).unwrap()],
5854                always_confirm: vec![],
5855                invalid_patterns: vec![],
5856            },
5857        );
5858        agent_settings::AgentSettings::override_global(settings, cx);
5859    });
5860
5861    #[allow(clippy::arc_with_non_send_sync)]
5862    let tool = Arc::new(crate::MovePathTool::new(project));
5863    let (event_stream, _rx) = crate::ToolCallEventStream::test();
5864
5865    let task = cx.update(|cx| {
5866        tool.run(
5867            ToolInput::resolved(crate::MovePathToolInput {
5868                source_path: "root/safe.txt".to_string(),
5869                destination_path: "root/protected/safe.txt".to_string(),
5870            }),
5871            event_stream,
5872            cx,
5873        )
5874    });
5875
5876    let result = task.await;
5877    assert!(
5878        result.is_err(),
5879        "expected move to be blocked due to destination path"
5880    );
5881    assert!(
5882        result.unwrap_err().contains("blocked"),
5883        "error should mention the move was blocked"
5884    );
5885}
5886
5887#[gpui::test]
5888async fn test_move_path_tool_denies_if_source_denied(cx: &mut TestAppContext) {
5889    init_test(cx);
5890
5891    let fs = FakeFs::new(cx.executor());
5892    fs.insert_tree(
5893        "/root",
5894        json!({
5895            "secret.txt": "secret content",
5896            "public": {}
5897        }),
5898    )
5899    .await;
5900    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
5901
5902    cx.update(|cx| {
5903        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
5904        settings.tool_permissions.tools.insert(
5905            MovePathTool::NAME.into(),
5906            agent_settings::ToolRules {
5907                default: Some(settings::ToolPermissionMode::Allow),
5908                always_allow: vec![],
5909                always_deny: vec![agent_settings::CompiledRegex::new(r"secret", false).unwrap()],
5910                always_confirm: vec![],
5911                invalid_patterns: vec![],
5912            },
5913        );
5914        agent_settings::AgentSettings::override_global(settings, cx);
5915    });
5916
5917    #[allow(clippy::arc_with_non_send_sync)]
5918    let tool = Arc::new(crate::MovePathTool::new(project));
5919    let (event_stream, _rx) = crate::ToolCallEventStream::test();
5920
5921    let task = cx.update(|cx| {
5922        tool.run(
5923            ToolInput::resolved(crate::MovePathToolInput {
5924                source_path: "root/secret.txt".to_string(),
5925                destination_path: "root/public/not_secret.txt".to_string(),
5926            }),
5927            event_stream,
5928            cx,
5929        )
5930    });
5931
5932    let result = task.await;
5933    assert!(
5934        result.is_err(),
5935        "expected move to be blocked due to source path"
5936    );
5937    assert!(
5938        result.unwrap_err().contains("blocked"),
5939        "error should mention the move was blocked"
5940    );
5941}
5942
5943#[gpui::test]
5944async fn test_copy_path_tool_deny_rule_blocks_copy(cx: &mut TestAppContext) {
5945    init_test(cx);
5946
5947    let fs = FakeFs::new(cx.executor());
5948    fs.insert_tree(
5949        "/root",
5950        json!({
5951            "confidential.txt": "confidential data",
5952            "dest": {}
5953        }),
5954    )
5955    .await;
5956    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
5957
5958    cx.update(|cx| {
5959        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
5960        settings.tool_permissions.tools.insert(
5961            CopyPathTool::NAME.into(),
5962            agent_settings::ToolRules {
5963                default: Some(settings::ToolPermissionMode::Allow),
5964                always_allow: vec![],
5965                always_deny: vec![
5966                    agent_settings::CompiledRegex::new(r"confidential", false).unwrap(),
5967                ],
5968                always_confirm: vec![],
5969                invalid_patterns: vec![],
5970            },
5971        );
5972        agent_settings::AgentSettings::override_global(settings, cx);
5973    });
5974
5975    #[allow(clippy::arc_with_non_send_sync)]
5976    let tool = Arc::new(crate::CopyPathTool::new(project));
5977    let (event_stream, _rx) = crate::ToolCallEventStream::test();
5978
5979    let task = cx.update(|cx| {
5980        tool.run(
5981            ToolInput::resolved(crate::CopyPathToolInput {
5982                source_path: "root/confidential.txt".to_string(),
5983                destination_path: "root/dest/copy.txt".to_string(),
5984            }),
5985            event_stream,
5986            cx,
5987        )
5988    });
5989
5990    let result = task.await;
5991    assert!(result.is_err(), "expected copy to be blocked");
5992    assert!(
5993        result.unwrap_err().contains("blocked"),
5994        "error should mention the copy was blocked"
5995    );
5996}
5997
5998#[gpui::test]
5999async fn test_save_file_tool_denies_if_any_path_denied(cx: &mut TestAppContext) {
6000    init_test(cx);
6001
6002    let fs = FakeFs::new(cx.executor());
6003    fs.insert_tree(
6004        "/root",
6005        json!({
6006            "normal.txt": "normal content",
6007            "readonly": {
6008                "config.txt": "readonly content"
6009            }
6010        }),
6011    )
6012    .await;
6013    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
6014
6015    cx.update(|cx| {
6016        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6017        settings.tool_permissions.tools.insert(
6018            SaveFileTool::NAME.into(),
6019            agent_settings::ToolRules {
6020                default: Some(settings::ToolPermissionMode::Allow),
6021                always_allow: vec![],
6022                always_deny: vec![agent_settings::CompiledRegex::new(r"readonly", false).unwrap()],
6023                always_confirm: vec![],
6024                invalid_patterns: vec![],
6025            },
6026        );
6027        agent_settings::AgentSettings::override_global(settings, cx);
6028    });
6029
6030    #[allow(clippy::arc_with_non_send_sync)]
6031    let tool = Arc::new(crate::SaveFileTool::new(project));
6032    let (event_stream, _rx) = crate::ToolCallEventStream::test();
6033
6034    let task = cx.update(|cx| {
6035        tool.run(
6036            ToolInput::resolved(crate::SaveFileToolInput {
6037                paths: vec![
6038                    std::path::PathBuf::from("root/normal.txt"),
6039                    std::path::PathBuf::from("root/readonly/config.txt"),
6040                ],
6041            }),
6042            event_stream,
6043            cx,
6044        )
6045    });
6046
6047    let result = task.await;
6048    assert!(
6049        result.is_err(),
6050        "expected save to be blocked due to denied path"
6051    );
6052    assert!(
6053        result.unwrap_err().contains("blocked"),
6054        "error should mention the save was blocked"
6055    );
6056}
6057
6058#[gpui::test]
6059async fn test_save_file_tool_respects_deny_rules(cx: &mut TestAppContext) {
6060    init_test(cx);
6061
6062    let fs = FakeFs::new(cx.executor());
6063    fs.insert_tree("/root", json!({"config.secret": "secret config"}))
6064        .await;
6065    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
6066
6067    cx.update(|cx| {
6068        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6069        settings.tool_permissions.tools.insert(
6070            SaveFileTool::NAME.into(),
6071            agent_settings::ToolRules {
6072                default: Some(settings::ToolPermissionMode::Allow),
6073                always_allow: vec![],
6074                always_deny: vec![agent_settings::CompiledRegex::new(r"\.secret$", false).unwrap()],
6075                always_confirm: vec![],
6076                invalid_patterns: vec![],
6077            },
6078        );
6079        agent_settings::AgentSettings::override_global(settings, cx);
6080    });
6081
6082    #[allow(clippy::arc_with_non_send_sync)]
6083    let tool = Arc::new(crate::SaveFileTool::new(project));
6084    let (event_stream, _rx) = crate::ToolCallEventStream::test();
6085
6086    let task = cx.update(|cx| {
6087        tool.run(
6088            ToolInput::resolved(crate::SaveFileToolInput {
6089                paths: vec![std::path::PathBuf::from("root/config.secret")],
6090            }),
6091            event_stream,
6092            cx,
6093        )
6094    });
6095
6096    let result = task.await;
6097    assert!(result.is_err(), "expected save to be blocked");
6098    assert!(
6099        result.unwrap_err().contains("blocked"),
6100        "error should mention the save was blocked"
6101    );
6102}
6103
6104#[gpui::test]
6105async fn test_web_search_tool_deny_rule_blocks_search(cx: &mut TestAppContext) {
6106    init_test(cx);
6107
6108    cx.update(|cx| {
6109        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6110        settings.tool_permissions.tools.insert(
6111            WebSearchTool::NAME.into(),
6112            agent_settings::ToolRules {
6113                default: Some(settings::ToolPermissionMode::Allow),
6114                always_allow: vec![],
6115                always_deny: vec![
6116                    agent_settings::CompiledRegex::new(r"internal\.company", false).unwrap(),
6117                ],
6118                always_confirm: vec![],
6119                invalid_patterns: vec![],
6120            },
6121        );
6122        agent_settings::AgentSettings::override_global(settings, cx);
6123    });
6124
6125    #[allow(clippy::arc_with_non_send_sync)]
6126    let tool = Arc::new(crate::WebSearchTool);
6127    let (event_stream, _rx) = crate::ToolCallEventStream::test();
6128
6129    let input: crate::WebSearchToolInput =
6130        serde_json::from_value(json!({"query": "internal.company.com secrets"})).unwrap();
6131
6132    let task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));
6133
6134    let result = task.await;
6135    assert!(result.is_err(), "expected search to be blocked");
6136    match result.unwrap_err() {
6137        crate::WebSearchToolOutput::Error { error } => {
6138            assert!(
6139                error.contains("blocked"),
6140                "error should mention the search was blocked"
6141            );
6142        }
6143        other => panic!("expected Error variant, got: {other:?}"),
6144    }
6145}
6146
6147#[gpui::test]
6148async fn test_edit_file_tool_allow_rule_skips_confirmation(cx: &mut TestAppContext) {
6149    init_test(cx);
6150
6151    let fs = FakeFs::new(cx.executor());
6152    fs.insert_tree("/root", json!({"README.md": "# Hello"}))
6153        .await;
6154    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
6155
6156    cx.update(|cx| {
6157        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6158        settings.tool_permissions.tools.insert(
6159            EditFileTool::NAME.into(),
6160            agent_settings::ToolRules {
6161                default: Some(settings::ToolPermissionMode::Confirm),
6162                always_allow: vec![agent_settings::CompiledRegex::new(r"\.md$", false).unwrap()],
6163                always_deny: vec![],
6164                always_confirm: vec![],
6165                invalid_patterns: vec![],
6166            },
6167        );
6168        agent_settings::AgentSettings::override_global(settings, cx);
6169    });
6170
6171    let context_server_registry =
6172        cx.new(|cx| crate::ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
6173    let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
6174    let templates = crate::Templates::new();
6175    let thread = cx.new(|cx| {
6176        crate::Thread::new(
6177            project.clone(),
6178            cx.new(|_cx| prompt_store::ProjectContext::default()),
6179            context_server_registry,
6180            templates.clone(),
6181            None,
6182            cx,
6183        )
6184    });
6185
6186    #[allow(clippy::arc_with_non_send_sync)]
6187    let tool = Arc::new(crate::EditFileTool::new(
6188        project,
6189        thread.downgrade(),
6190        language_registry,
6191        templates,
6192    ));
6193    let (event_stream, mut rx) = crate::ToolCallEventStream::test();
6194
6195    let _task = cx.update(|cx| {
6196        tool.run(
6197            ToolInput::resolved(crate::EditFileToolInput {
6198                display_description: "Edit README".to_string(),
6199                path: "root/README.md".into(),
6200                mode: crate::EditFileMode::Edit,
6201            }),
6202            event_stream,
6203            cx,
6204        )
6205    });
6206
6207    cx.run_until_parked();
6208
6209    let event = rx.try_recv();
6210    assert!(
6211        !matches!(event, Ok(Ok(ThreadEvent::ToolCallAuthorization(_)))),
6212        "expected no authorization request for allowed .md file"
6213    );
6214}
6215
6216#[gpui::test]
6217async fn test_edit_file_tool_allow_still_prompts_for_local_settings(cx: &mut TestAppContext) {
6218    init_test(cx);
6219
6220    let fs = FakeFs::new(cx.executor());
6221    fs.insert_tree(
6222        "/root",
6223        json!({
6224            ".zed": {
6225                "settings.json": "{}"
6226            },
6227            "README.md": "# Hello"
6228        }),
6229    )
6230    .await;
6231    let project = Project::test(fs.clone(), ["/root".as_ref()], cx).await;
6232
6233    cx.update(|cx| {
6234        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6235        settings.tool_permissions.default = settings::ToolPermissionMode::Allow;
6236        agent_settings::AgentSettings::override_global(settings, cx);
6237    });
6238
6239    let context_server_registry =
6240        cx.new(|cx| crate::ContextServerRegistry::new(project.read(cx).context_server_store(), cx));
6241    let language_registry = project.read_with(cx, |project, _cx| project.languages().clone());
6242    let templates = crate::Templates::new();
6243    let thread = cx.new(|cx| {
6244        crate::Thread::new(
6245            project.clone(),
6246            cx.new(|_cx| prompt_store::ProjectContext::default()),
6247            context_server_registry,
6248            templates.clone(),
6249            None,
6250            cx,
6251        )
6252    });
6253
6254    #[allow(clippy::arc_with_non_send_sync)]
6255    let tool = Arc::new(crate::EditFileTool::new(
6256        project,
6257        thread.downgrade(),
6258        language_registry,
6259        templates,
6260    ));
6261
6262    // Editing a file inside .zed/ should still prompt even with global default: allow,
6263    // because local settings paths are sensitive and require confirmation regardless.
6264    let (event_stream, mut rx) = crate::ToolCallEventStream::test();
6265    let _task = cx.update(|cx| {
6266        tool.run(
6267            ToolInput::resolved(crate::EditFileToolInput {
6268                display_description: "Edit local settings".to_string(),
6269                path: "root/.zed/settings.json".into(),
6270                mode: crate::EditFileMode::Edit,
6271            }),
6272            event_stream,
6273            cx,
6274        )
6275    });
6276
6277    let _update = rx.expect_update_fields().await;
6278    let _auth = rx.expect_authorization().await;
6279}
6280
6281#[gpui::test]
6282async fn test_fetch_tool_deny_rule_blocks_url(cx: &mut TestAppContext) {
6283    init_test(cx);
6284
6285    cx.update(|cx| {
6286        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6287        settings.tool_permissions.tools.insert(
6288            FetchTool::NAME.into(),
6289            agent_settings::ToolRules {
6290                default: Some(settings::ToolPermissionMode::Allow),
6291                always_allow: vec![],
6292                always_deny: vec![
6293                    agent_settings::CompiledRegex::new(r"internal\.company\.com", false).unwrap(),
6294                ],
6295                always_confirm: vec![],
6296                invalid_patterns: vec![],
6297            },
6298        );
6299        agent_settings::AgentSettings::override_global(settings, cx);
6300    });
6301
6302    let http_client = gpui::http_client::FakeHttpClient::with_200_response();
6303
6304    #[allow(clippy::arc_with_non_send_sync)]
6305    let tool = Arc::new(crate::FetchTool::new(http_client));
6306    let (event_stream, _rx) = crate::ToolCallEventStream::test();
6307
6308    let input: crate::FetchToolInput =
6309        serde_json::from_value(json!({"url": "https://internal.company.com/api"})).unwrap();
6310
6311    let task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));
6312
6313    let result = task.await;
6314    assert!(result.is_err(), "expected fetch to be blocked");
6315    assert!(
6316        result.unwrap_err().contains("blocked"),
6317        "error should mention the fetch was blocked"
6318    );
6319}
6320
6321#[gpui::test]
6322async fn test_fetch_tool_allow_rule_skips_confirmation(cx: &mut TestAppContext) {
6323    init_test(cx);
6324
6325    cx.update(|cx| {
6326        let mut settings = agent_settings::AgentSettings::get_global(cx).clone();
6327        settings.tool_permissions.tools.insert(
6328            FetchTool::NAME.into(),
6329            agent_settings::ToolRules {
6330                default: Some(settings::ToolPermissionMode::Confirm),
6331                always_allow: vec![agent_settings::CompiledRegex::new(r"docs\.rs", false).unwrap()],
6332                always_deny: vec![],
6333                always_confirm: vec![],
6334                invalid_patterns: vec![],
6335            },
6336        );
6337        agent_settings::AgentSettings::override_global(settings, cx);
6338    });
6339
6340    let http_client = gpui::http_client::FakeHttpClient::with_200_response();
6341
6342    #[allow(clippy::arc_with_non_send_sync)]
6343    let tool = Arc::new(crate::FetchTool::new(http_client));
6344    let (event_stream, mut rx) = crate::ToolCallEventStream::test();
6345
6346    let input: crate::FetchToolInput =
6347        serde_json::from_value(json!({"url": "https://docs.rs/some-crate"})).unwrap();
6348
6349    let _task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));
6350
6351    cx.run_until_parked();
6352
6353    let event = rx.try_recv();
6354    assert!(
6355        !matches!(event, Ok(Ok(ThreadEvent::ToolCallAuthorization(_)))),
6356        "expected no authorization request for allowed docs.rs URL"
6357    );
6358}
6359
6360#[gpui::test]
6361async fn test_queued_message_ends_turn_at_boundary(cx: &mut TestAppContext) {
6362    init_test(cx);
6363    always_allow_tools(cx);
6364
6365    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
6366    let fake_model = model.as_fake();
6367
6368    // Add a tool so we can simulate tool calls
6369    thread.update(cx, |thread, _cx| {
6370        thread.add_tool(EchoTool);
6371    });
6372
6373    // Start a turn by sending a message
6374    let mut events = thread
6375        .update(cx, |thread, cx| {
6376            thread.send(UserMessageId::new(), ["Use the echo tool"], cx)
6377        })
6378        .unwrap();
6379    cx.run_until_parked();
6380
6381    // Simulate the model making a tool call
6382    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
6383        LanguageModelToolUse {
6384            id: "tool_1".into(),
6385            name: "echo".into(),
6386            raw_input: r#"{"text": "hello"}"#.into(),
6387            input: json!({"text": "hello"}),
6388            is_input_complete: true,
6389            thought_signature: None,
6390        },
6391    ));
6392    fake_model
6393        .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::ToolUse));
6394
6395    // Signal that a message is queued before ending the stream
6396    thread.update(cx, |thread, _cx| {
6397        thread.set_has_queued_message(true);
6398    });
6399
6400    // Now end the stream - tool will run, and the boundary check should see the queue
6401    fake_model.end_last_completion_stream();
6402
6403    // Collect all events until the turn stops
6404    let all_events = collect_events_until_stop(&mut events, cx).await;
6405
6406    // Verify we received the tool call event
6407    let tool_call_ids: Vec<_> = all_events
6408        .iter()
6409        .filter_map(|e| match e {
6410            Ok(ThreadEvent::ToolCall(tc)) => Some(tc.tool_call_id.to_string()),
6411            _ => None,
6412        })
6413        .collect();
6414    assert_eq!(
6415        tool_call_ids,
6416        vec!["tool_1"],
6417        "Should have received a tool call event for our echo tool"
6418    );
6419
6420    // The turn should have stopped with EndTurn
6421    let stop_reasons = stop_events(all_events);
6422    assert_eq!(
6423        stop_reasons,
6424        vec![acp::StopReason::EndTurn],
6425        "Turn should have ended after tool completion due to queued message"
6426    );
6427
6428    // Verify the queued message flag is still set
6429    thread.update(cx, |thread, _cx| {
6430        assert!(
6431            thread.has_queued_message(),
6432            "Should still have queued message flag set"
6433        );
6434    });
6435
6436    // Thread should be idle now
6437    thread.update(cx, |thread, _cx| {
6438        assert!(
6439            thread.is_turn_complete(),
6440            "Thread should not be running after turn ends"
6441        );
6442    });
6443}
6444
6445#[gpui::test]
6446async fn test_streaming_tool_error_breaks_stream_loop_immediately(cx: &mut TestAppContext) {
6447    init_test(cx);
6448    always_allow_tools(cx);
6449
6450    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
6451    let fake_model = model.as_fake();
6452
6453    thread.update(cx, |thread, _cx| {
6454        thread.add_tool(StreamingFailingEchoTool {
6455            receive_chunks_until_failure: 1,
6456        });
6457    });
6458
6459    let _events = thread
6460        .update(cx, |thread, cx| {
6461            thread.send(
6462                UserMessageId::new(),
6463                ["Use the streaming_failing_echo tool"],
6464                cx,
6465            )
6466        })
6467        .unwrap();
6468    cx.run_until_parked();
6469
6470    let tool_use = LanguageModelToolUse {
6471        id: "call_1".into(),
6472        name: StreamingFailingEchoTool::NAME.into(),
6473        raw_input: "hello".into(),
6474        input: json!({}),
6475        is_input_complete: false,
6476        thought_signature: None,
6477    };
6478
6479    fake_model
6480        .send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(tool_use.clone()));
6481
6482    cx.run_until_parked();
6483
6484    let completions = fake_model.pending_completions();
6485    let last_completion = completions.last().unwrap();
6486
6487    assert_eq!(
6488        last_completion.messages[1..],
6489        vec![
6490            LanguageModelRequestMessage {
6491                role: Role::User,
6492                content: vec!["Use the streaming_failing_echo tool".into()],
6493                cache: false,
6494                reasoning_details: None,
6495            },
6496            LanguageModelRequestMessage {
6497                role: Role::Assistant,
6498                content: vec![language_model::MessageContent::ToolUse(tool_use.clone())],
6499                cache: false,
6500                reasoning_details: None,
6501            },
6502            LanguageModelRequestMessage {
6503                role: Role::User,
6504                content: vec![language_model::MessageContent::ToolResult(
6505                    LanguageModelToolResult {
6506                        tool_use_id: tool_use.id.clone(),
6507                        tool_name: tool_use.name,
6508                        is_error: true,
6509                        content: "failed".into(),
6510                        output: Some("failed".into()),
6511                    }
6512                )],
6513                cache: true,
6514                reasoning_details: None,
6515            },
6516        ]
6517    );
6518}
6519
6520#[gpui::test]
6521async fn test_streaming_tool_error_waits_for_prior_tools_to_complete(cx: &mut TestAppContext) {
6522    init_test(cx);
6523    always_allow_tools(cx);
6524
6525    let ThreadTest { model, thread, .. } = setup(cx, TestModel::Fake).await;
6526    let fake_model = model.as_fake();
6527
6528    let (complete_streaming_echo_tool_call_tx, complete_streaming_echo_tool_call_rx) =
6529        oneshot::channel();
6530
6531    thread.update(cx, |thread, _cx| {
6532        thread.add_tool(
6533            StreamingEchoTool::new().with_wait_until_complete(complete_streaming_echo_tool_call_rx),
6534        );
6535        thread.add_tool(StreamingFailingEchoTool {
6536            receive_chunks_until_failure: 1,
6537        });
6538    });
6539
6540    let _events = thread
6541        .update(cx, |thread, cx| {
6542            thread.send(
6543                UserMessageId::new(),
6544                ["Use the streaming_echo tool and the streaming_failing_echo tool"],
6545                cx,
6546            )
6547        })
6548        .unwrap();
6549    cx.run_until_parked();
6550
6551    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
6552        LanguageModelToolUse {
6553            id: "call_1".into(),
6554            name: StreamingEchoTool::NAME.into(),
6555            raw_input: "hello".into(),
6556            input: json!({ "text": "hello" }),
6557            is_input_complete: false,
6558            thought_signature: None,
6559        },
6560    ));
6561    let first_tool_use = LanguageModelToolUse {
6562        id: "call_1".into(),
6563        name: StreamingEchoTool::NAME.into(),
6564        raw_input: "hello world".into(),
6565        input: json!({ "text": "hello world" }),
6566        is_input_complete: true,
6567        thought_signature: None,
6568    };
6569    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
6570        first_tool_use.clone(),
6571    ));
6572    let second_tool_use = LanguageModelToolUse {
6573        name: StreamingFailingEchoTool::NAME.into(),
6574        raw_input: "hello".into(),
6575        input: json!({ "text": "hello" }),
6576        is_input_complete: false,
6577        thought_signature: None,
6578        id: "call_2".into(),
6579    };
6580    fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
6581        second_tool_use.clone(),
6582    ));
6583
6584    cx.run_until_parked();
6585
6586    complete_streaming_echo_tool_call_tx.send(()).unwrap();
6587
6588    cx.run_until_parked();
6589
6590    let completions = fake_model.pending_completions();
6591    let last_completion = completions.last().unwrap();
6592
6593    assert_eq!(
6594        last_completion.messages[1..],
6595        vec![
6596            LanguageModelRequestMessage {
6597                role: Role::User,
6598                content: vec![
6599                    "Use the streaming_echo tool and the streaming_failing_echo tool".into()
6600                ],
6601                cache: false,
6602                reasoning_details: None,
6603            },
6604            LanguageModelRequestMessage {
6605                role: Role::Assistant,
6606                content: vec![
6607                    language_model::MessageContent::ToolUse(first_tool_use.clone()),
6608                    language_model::MessageContent::ToolUse(second_tool_use.clone())
6609                ],
6610                cache: false,
6611                reasoning_details: None,
6612            },
6613            LanguageModelRequestMessage {
6614                role: Role::User,
6615                content: vec![
6616                    language_model::MessageContent::ToolResult(LanguageModelToolResult {
6617                        tool_use_id: second_tool_use.id.clone(),
6618                        tool_name: second_tool_use.name,
6619                        is_error: true,
6620                        content: "failed".into(),
6621                        output: Some("failed".into()),
6622                    }),
6623                    language_model::MessageContent::ToolResult(LanguageModelToolResult {
6624                        tool_use_id: first_tool_use.id.clone(),
6625                        tool_name: first_tool_use.name,
6626                        is_error: false,
6627                        content: "hello world".into(),
6628                        output: Some("hello world".into()),
6629                    }),
6630                ],
6631                cache: true,
6632                reasoning_details: None,
6633            },
6634        ]
6635    );
6636}
6637
6638#[gpui::test]
6639async fn test_mid_turn_model_and_settings_refresh(cx: &mut TestAppContext) {
6640    let ThreadTest {
6641        model, thread, fs, ..
6642    } = setup(cx, TestModel::Fake).await;
6643    let fake_model_a = model.as_fake();
6644
6645    thread.update(cx, |thread, _cx| {
6646        thread.add_tool(EchoTool);
6647        thread.add_tool(DelayTool);
6648    });
6649
6650    // Set up two profiles: profile-a has both tools, profile-b has only DelayTool.
6651    fs.insert_file(
6652        paths::settings_file(),
6653        json!({
6654            "agent": {
6655                "profiles": {
6656                    "profile-a": {
6657                        "name": "Profile A",
6658                        "tools": {
6659                            EchoTool::NAME: true,
6660                            DelayTool::NAME: true,
6661                        }
6662                    },
6663                    "profile-b": {
6664                        "name": "Profile B",
6665                        "tools": {
6666                            DelayTool::NAME: true,
6667                        }
6668                    }
6669                }
6670            }
6671        })
6672        .to_string()
6673        .into_bytes(),
6674    )
6675    .await;
6676    cx.run_until_parked();
6677
6678    thread.update(cx, |thread, cx| {
6679        thread.set_profile(AgentProfileId("profile-a".into()), cx);
6680        thread.set_thinking_enabled(false, cx);
6681    });
6682
6683    // Send a message — first iteration starts with model A, profile-a, thinking off.
6684    thread
6685        .update(cx, |thread, cx| {
6686            thread.send(UserMessageId::new(), ["test mid-turn refresh"], cx)
6687        })
6688        .unwrap();
6689    cx.run_until_parked();
6690
6691    // Verify first request has both tools and thinking disabled.
6692    let completions = fake_model_a.pending_completions();
6693    assert_eq!(completions.len(), 1);
6694    let first_tools = tool_names_for_completion(&completions[0]);
6695    assert_eq!(first_tools, vec![DelayTool::NAME, EchoTool::NAME]);
6696    assert!(!completions[0].thinking_allowed);
6697
6698    // Model A responds with an echo tool call.
6699    fake_model_a.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
6700        LanguageModelToolUse {
6701            id: "tool_1".into(),
6702            name: "echo".into(),
6703            raw_input: r#"{"text":"hello"}"#.into(),
6704            input: json!({"text": "hello"}),
6705            is_input_complete: true,
6706            thought_signature: None,
6707        },
6708    ));
6709    fake_model_a.end_last_completion_stream();
6710
6711    // Before the next iteration runs, switch to profile-b (only DelayTool),
6712    // swap in a new model, and enable thinking.
6713    let fake_model_b = Arc::new(FakeLanguageModel::with_id_and_thinking(
6714        "test-provider",
6715        "model-b",
6716        "Model B",
6717        true,
6718    ));
6719    thread.update(cx, |thread, cx| {
6720        thread.set_profile(AgentProfileId("profile-b".into()), cx);
6721        thread.set_model(fake_model_b.clone() as Arc<dyn LanguageModel>, cx);
6722        thread.set_thinking_enabled(true, cx);
6723    });
6724
6725    // Run until parked — processes the echo tool call, loops back, picks up
6726    // the new model/profile/thinking, and makes a second request to model B.
6727    cx.run_until_parked();
6728
6729    // The second request should have gone to model B.
6730    let model_b_completions = fake_model_b.pending_completions();
6731    assert_eq!(
6732        model_b_completions.len(),
6733        1,
6734        "second request should go to model B"
6735    );
6736
6737    // Profile-b only has DelayTool, so echo should be gone.
6738    let second_tools = tool_names_for_completion(&model_b_completions[0]);
6739    assert_eq!(second_tools, vec![DelayTool::NAME]);
6740
6741    // Thinking should now be enabled.
6742    assert!(model_b_completions[0].thinking_allowed);
6743}