agent.rs

   1mod db;
   2mod edit_agent;
   3mod legacy_thread;
   4mod native_agent_server;
   5pub mod outline;
   6mod pattern_extraction;
   7mod templates;
   8#[cfg(test)]
   9mod tests;
  10mod thread;
  11mod thread_store;
  12mod tool_permissions;
  13mod tools;
  14
  15use context_server::ContextServerId;
  16pub use db::*;
  17use itertools::Itertools;
  18pub use native_agent_server::NativeAgentServer;
  19pub use pattern_extraction::*;
  20pub use shell_command_parser::extract_commands;
  21pub use templates::*;
  22pub use thread::*;
  23pub use thread_store::*;
  24pub use tool_permissions::*;
  25pub use tools::*;
  26
  27use acp_thread::{
  28    AcpThread, AgentModelSelector, AgentSessionInfo, AgentSessionList, AgentSessionListRequest,
  29    AgentSessionListResponse, TokenUsageRatio, UserMessageId,
  30};
  31use agent_client_protocol as acp;
  32use anyhow::{Context as _, Result, anyhow};
  33use chrono::{DateTime, Utc};
  34use collections::{HashMap, HashSet, IndexMap};
  35use fs::Fs;
  36use futures::channel::{mpsc, oneshot};
  37use futures::future::Shared;
  38use futures::{FutureExt as _, StreamExt as _, future};
  39use gpui::{
  40    App, AppContext, AsyncApp, Context, Entity, EntityId, SharedString, Subscription, Task,
  41    WeakEntity,
  42};
  43use language_model::{IconOrSvg, LanguageModel, LanguageModelProvider, LanguageModelRegistry};
  44use project::{AgentId, Project, ProjectItem, ProjectPath, Worktree};
  45use prompt_store::{
  46    ProjectContext, PromptStore, RULES_FILE_NAMES, RulesFileContext, UserRulesContext,
  47    WorktreeContext,
  48};
  49use serde::{Deserialize, Serialize};
  50use settings::{LanguageModelSelection, update_settings_file};
  51use std::any::Any;
  52use std::path::PathBuf;
  53use std::rc::Rc;
  54use std::sync::{Arc, LazyLock};
  55use util::ResultExt;
  56use util::path_list::PathList;
  57use util::rel_path::RelPath;
  58
  59#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
  60pub struct ProjectSnapshot {
  61    pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
  62    pub timestamp: DateTime<Utc>,
  63}
  64
  65pub struct RulesLoadingError {
  66    pub message: SharedString,
  67}
  68
  69struct ProjectState {
  70    project: Entity<Project>,
  71    project_context: Entity<ProjectContext>,
  72    project_context_needs_refresh: watch::Sender<()>,
  73    _maintain_project_context: Task<Result<()>>,
  74    context_server_registry: Entity<ContextServerRegistry>,
  75    _subscriptions: Vec<Subscription>,
  76}
  77
  78/// Holds both the internal Thread and the AcpThread for a session
  79struct Session {
  80    /// The internal thread that processes messages
  81    thread: Entity<Thread>,
  82    /// The ACP thread that handles protocol communication
  83    acp_thread: Entity<acp_thread::AcpThread>,
  84    project_id: EntityId,
  85    pending_save: Task<()>,
  86    _subscriptions: Vec<Subscription>,
  87}
  88
  89pub struct LanguageModels {
  90    /// Access language model by ID
  91    models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
  92    /// Cached list for returning language model information
  93    model_list: acp_thread::AgentModelList,
  94    refresh_models_rx: watch::Receiver<()>,
  95    refresh_models_tx: watch::Sender<()>,
  96    _authenticate_all_providers_task: Task<()>,
  97}
  98
  99impl LanguageModels {
 100    fn new(cx: &mut App) -> Self {
 101        let (refresh_models_tx, refresh_models_rx) = watch::channel(());
 102
 103        let mut this = Self {
 104            models: HashMap::default(),
 105            model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
 106            refresh_models_rx,
 107            refresh_models_tx,
 108            _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
 109        };
 110        this.refresh_list(cx);
 111        this
 112    }
 113
 114    fn refresh_list(&mut self, cx: &App) {
 115        let providers = LanguageModelRegistry::global(cx)
 116            .read(cx)
 117            .visible_providers()
 118            .into_iter()
 119            .filter(|provider| provider.is_authenticated(cx))
 120            .collect::<Vec<_>>();
 121
 122        let mut language_model_list = IndexMap::default();
 123        let mut recommended_models = HashSet::default();
 124
 125        let mut recommended = Vec::new();
 126        for provider in &providers {
 127            for model in provider.recommended_models(cx) {
 128                recommended_models.insert((model.provider_id(), model.id()));
 129                recommended.push(Self::map_language_model_to_info(&model, provider));
 130            }
 131        }
 132        if !recommended.is_empty() {
 133            language_model_list.insert(
 134                acp_thread::AgentModelGroupName("Recommended".into()),
 135                recommended,
 136            );
 137        }
 138
 139        let mut models = HashMap::default();
 140        for provider in providers {
 141            let mut provider_models = Vec::new();
 142            for model in provider.provided_models(cx) {
 143                let model_info = Self::map_language_model_to_info(&model, &provider);
 144                let model_id = model_info.id.clone();
 145                provider_models.push(model_info);
 146                models.insert(model_id, model);
 147            }
 148            if !provider_models.is_empty() {
 149                language_model_list.insert(
 150                    acp_thread::AgentModelGroupName(provider.name().0.clone()),
 151                    provider_models,
 152                );
 153            }
 154        }
 155
 156        self.models = models;
 157        self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
 158        self.refresh_models_tx.send(()).ok();
 159    }
 160
 161    fn watch(&self) -> watch::Receiver<()> {
 162        self.refresh_models_rx.clone()
 163    }
 164
 165    pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
 166        self.models.get(model_id).cloned()
 167    }
 168
 169    fn map_language_model_to_info(
 170        model: &Arc<dyn LanguageModel>,
 171        provider: &Arc<dyn LanguageModelProvider>,
 172    ) -> acp_thread::AgentModelInfo {
 173        acp_thread::AgentModelInfo {
 174            id: Self::model_id(model),
 175            name: model.name().0,
 176            description: None,
 177            icon: Some(match provider.icon() {
 178                IconOrSvg::Svg(path) => acp_thread::AgentModelIcon::Path(path),
 179                IconOrSvg::Icon(name) => acp_thread::AgentModelIcon::Named(name),
 180            }),
 181            is_latest: model.is_latest(),
 182            cost: model.model_cost_info().map(|cost| cost.to_shared_string()),
 183        }
 184    }
 185
 186    fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
 187        acp::ModelId::new(format!("{}/{}", model.provider_id().0, model.id().0))
 188    }
 189
 190    fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
 191        let authenticate_all_providers = LanguageModelRegistry::global(cx)
 192            .read(cx)
 193            .visible_providers()
 194            .iter()
 195            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
 196            .collect::<Vec<_>>();
 197
 198        cx.background_spawn(async move {
 199            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
 200                if let Err(err) = authenticate_task.await {
 201                    match err {
 202                        language_model::AuthenticateError::CredentialsNotFound => {
 203                            // Since we're authenticating these providers in the
 204                            // background for the purposes of populating the
 205                            // language selector, we don't care about providers
 206                            // where the credentials are not found.
 207                        }
 208                        language_model::AuthenticateError::ConnectionRefused => {
 209                            // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
 210                            // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
 211                            // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
 212                        }
 213                        _ => {
 214                            // Some providers have noisy failure states that we
 215                            // don't want to spam the logs with every time the
 216                            // language model selector is initialized.
 217                            //
 218                            // Ideally these should have more clear failure modes
 219                            // that we know are safe to ignore here, like what we do
 220                            // with `CredentialsNotFound` above.
 221                            match provider_id.0.as_ref() {
 222                                "lmstudio" | "ollama" => {
 223                                    // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
 224                                    //
 225                                    // These fail noisily, so we don't log them.
 226                                }
 227                                "copilot_chat" => {
 228                                    // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
 229                                }
 230                                _ => {
 231                                    log::error!(
 232                                        "Failed to authenticate provider: {}: {err:#}",
 233                                        provider_name.0
 234                                    );
 235                                }
 236                            }
 237                        }
 238                    }
 239                }
 240            }
 241        })
 242    }
 243}
 244
 245pub struct NativeAgent {
 246    /// Session ID -> Session mapping
 247    sessions: HashMap<acp::SessionId, Session>,
 248    thread_store: Entity<ThreadStore>,
 249    /// Project-specific state keyed by project EntityId
 250    projects: HashMap<EntityId, ProjectState>,
 251    /// Shared templates for all threads
 252    templates: Arc<Templates>,
 253    /// Cached model information
 254    models: LanguageModels,
 255    prompt_store: Option<Entity<PromptStore>>,
 256    fs: Arc<dyn Fs>,
 257    _subscriptions: Vec<Subscription>,
 258}
 259
 260impl NativeAgent {
 261    pub fn new(
 262        thread_store: Entity<ThreadStore>,
 263        templates: Arc<Templates>,
 264        prompt_store: Option<Entity<PromptStore>>,
 265        fs: Arc<dyn Fs>,
 266        cx: &mut App,
 267    ) -> Entity<NativeAgent> {
 268        log::debug!("Creating new NativeAgent");
 269
 270        cx.new(|cx| {
 271            let mut subscriptions = vec![cx.subscribe(
 272                &LanguageModelRegistry::global(cx),
 273                Self::handle_models_updated_event,
 274            )];
 275            if let Some(prompt_store) = prompt_store.as_ref() {
 276                subscriptions.push(cx.subscribe(prompt_store, Self::handle_prompts_updated_event))
 277            }
 278
 279            Self {
 280                sessions: HashMap::default(),
 281                thread_store,
 282                projects: HashMap::default(),
 283                templates,
 284                models: LanguageModels::new(cx),
 285                prompt_store,
 286                fs,
 287                _subscriptions: subscriptions,
 288            }
 289        })
 290    }
 291
 292    fn new_session(
 293        &mut self,
 294        project: Entity<Project>,
 295        cx: &mut Context<Self>,
 296    ) -> Entity<AcpThread> {
 297        let project_id = self.get_or_create_project_state(&project, cx);
 298        let project_state = &self.projects[&project_id];
 299
 300        let registry = LanguageModelRegistry::read_global(cx);
 301        let available_count = registry.available_models(cx).count();
 302        log::debug!("Total available models: {}", available_count);
 303
 304        let default_model = registry.default_model().and_then(|default_model| {
 305            self.models
 306                .model_from_id(&LanguageModels::model_id(&default_model.model))
 307        });
 308        let thread = cx.new(|cx| {
 309            Thread::new(
 310                project,
 311                project_state.project_context.clone(),
 312                project_state.context_server_registry.clone(),
 313                self.templates.clone(),
 314                default_model,
 315                cx,
 316            )
 317        });
 318
 319        self.register_session(thread, project_id, cx)
 320    }
 321
 322    fn register_session(
 323        &mut self,
 324        thread_handle: Entity<Thread>,
 325        project_id: EntityId,
 326        cx: &mut Context<Self>,
 327    ) -> Entity<AcpThread> {
 328        let connection = Rc::new(NativeAgentConnection(cx.entity()));
 329
 330        let thread = thread_handle.read(cx);
 331        let session_id = thread.id().clone();
 332        let parent_session_id = thread.parent_thread_id();
 333        let title = thread.title();
 334        let draft_prompt = thread.draft_prompt().map(Vec::from);
 335        let scroll_position = thread.ui_scroll_position();
 336        let token_usage = thread.latest_token_usage();
 337        let project = thread.project.clone();
 338        let action_log = thread.action_log.clone();
 339        let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
 340        let acp_thread = cx.new(|cx| {
 341            let mut acp_thread = acp_thread::AcpThread::new(
 342                parent_session_id,
 343                title,
 344                None,
 345                connection,
 346                project.clone(),
 347                action_log.clone(),
 348                session_id.clone(),
 349                prompt_capabilities_rx,
 350                cx,
 351            );
 352            acp_thread.set_draft_prompt(draft_prompt);
 353            acp_thread.set_ui_scroll_position(scroll_position);
 354            acp_thread.update_token_usage(token_usage, cx);
 355            acp_thread
 356        });
 357
 358        let registry = LanguageModelRegistry::read_global(cx);
 359        let summarization_model = registry.thread_summary_model().map(|c| c.model);
 360
 361        let weak = cx.weak_entity();
 362        let weak_thread = thread_handle.downgrade();
 363        thread_handle.update(cx, |thread, cx| {
 364            thread.set_summarization_model(summarization_model, cx);
 365            thread.add_default_tools(
 366                Rc::new(NativeThreadEnvironment {
 367                    acp_thread: acp_thread.downgrade(),
 368                    thread: weak_thread,
 369                    agent: weak,
 370                }) as _,
 371                cx,
 372            )
 373        });
 374
 375        let subscriptions = vec![
 376            cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
 377            cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
 378            cx.observe(&thread_handle, move |this, thread, cx| {
 379                this.save_thread(thread, cx)
 380            }),
 381        ];
 382
 383        self.sessions.insert(
 384            session_id,
 385            Session {
 386                thread: thread_handle,
 387                acp_thread: acp_thread.clone(),
 388                project_id,
 389                _subscriptions: subscriptions,
 390                pending_save: Task::ready(()),
 391            },
 392        );
 393
 394        self.update_available_commands_for_project(project_id, cx);
 395
 396        acp_thread
 397    }
 398
 399    pub fn models(&self) -> &LanguageModels {
 400        &self.models
 401    }
 402
 403    fn get_or_create_project_state(
 404        &mut self,
 405        project: &Entity<Project>,
 406        cx: &mut Context<Self>,
 407    ) -> EntityId {
 408        let project_id = project.entity_id();
 409        if self.projects.contains_key(&project_id) {
 410            return project_id;
 411        }
 412
 413        let project_context = cx.new(|_| ProjectContext::new(vec![], vec![]));
 414        self.register_project_with_initial_context(project.clone(), project_context, cx);
 415        if let Some(state) = self.projects.get_mut(&project_id) {
 416            state.project_context_needs_refresh.send(()).ok();
 417        }
 418        project_id
 419    }
 420
 421    fn register_project_with_initial_context(
 422        &mut self,
 423        project: Entity<Project>,
 424        project_context: Entity<ProjectContext>,
 425        cx: &mut Context<Self>,
 426    ) {
 427        let project_id = project.entity_id();
 428
 429        let context_server_store = project.read(cx).context_server_store();
 430        let context_server_registry =
 431            cx.new(|cx| ContextServerRegistry::new(context_server_store.clone(), cx));
 432
 433        let subscriptions = vec![
 434            cx.subscribe(&project, Self::handle_project_event),
 435            cx.subscribe(
 436                &context_server_store,
 437                Self::handle_context_server_store_updated,
 438            ),
 439            cx.subscribe(
 440                &context_server_registry,
 441                Self::handle_context_server_registry_event,
 442            ),
 443        ];
 444
 445        let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
 446            watch::channel(());
 447
 448        self.projects.insert(
 449            project_id,
 450            ProjectState {
 451                project,
 452                project_context,
 453                project_context_needs_refresh: project_context_needs_refresh_tx,
 454                _maintain_project_context: cx.spawn(async move |this, cx| {
 455                    Self::maintain_project_context(
 456                        this,
 457                        project_id,
 458                        project_context_needs_refresh_rx,
 459                        cx,
 460                    )
 461                    .await
 462                }),
 463                context_server_registry,
 464                _subscriptions: subscriptions,
 465            },
 466        );
 467    }
 468
 469    fn session_project_state(&self, session_id: &acp::SessionId) -> Option<&ProjectState> {
 470        self.sessions
 471            .get(session_id)
 472            .and_then(|session| self.projects.get(&session.project_id))
 473    }
 474
 475    async fn maintain_project_context(
 476        this: WeakEntity<Self>,
 477        project_id: EntityId,
 478        mut needs_refresh: watch::Receiver<()>,
 479        cx: &mut AsyncApp,
 480    ) -> Result<()> {
 481        while needs_refresh.changed().await.is_ok() {
 482            let project_context = this
 483                .update(cx, |this, cx| {
 484                    let state = this
 485                        .projects
 486                        .get(&project_id)
 487                        .context("project state not found")?;
 488                    anyhow::Ok(Self::build_project_context(
 489                        &state.project,
 490                        this.prompt_store.as_ref(),
 491                        cx,
 492                    ))
 493                })??
 494                .await;
 495            this.update(cx, |this, cx| {
 496                if let Some(state) = this.projects.get(&project_id) {
 497                    state
 498                        .project_context
 499                        .update(cx, |current_project_context, _cx| {
 500                            *current_project_context = project_context;
 501                        });
 502                }
 503            })?;
 504        }
 505
 506        Ok(())
 507    }
 508
 509    fn build_project_context(
 510        project: &Entity<Project>,
 511        prompt_store: Option<&Entity<PromptStore>>,
 512        cx: &mut App,
 513    ) -> Task<ProjectContext> {
 514        let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
 515        let worktree_tasks = worktrees
 516            .into_iter()
 517            .map(|worktree| {
 518                Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
 519            })
 520            .collect::<Vec<_>>();
 521        let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
 522            prompt_store.read_with(cx, |prompt_store, cx| {
 523                let prompts = prompt_store.default_prompt_metadata();
 524                let load_tasks = prompts.into_iter().map(|prompt_metadata| {
 525                    let contents = prompt_store.load(prompt_metadata.id, cx);
 526                    async move { (contents.await, prompt_metadata) }
 527                });
 528                cx.background_spawn(future::join_all(load_tasks))
 529            })
 530        } else {
 531            Task::ready(vec![])
 532        };
 533
 534        cx.spawn(async move |_cx| {
 535            let (worktrees, default_user_rules) =
 536                future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
 537
 538            let worktrees = worktrees
 539                .into_iter()
 540                .map(|(worktree, _rules_error)| {
 541                    // TODO: show error message
 542                    // if let Some(rules_error) = rules_error {
 543                    //     this.update(cx, |_, cx| cx.emit(rules_error)).ok();
 544                    // }
 545                    worktree
 546                })
 547                .collect::<Vec<_>>();
 548
 549            let default_user_rules = default_user_rules
 550                .into_iter()
 551                .flat_map(|(contents, prompt_metadata)| match contents {
 552                    Ok(contents) => Some(UserRulesContext {
 553                        uuid: prompt_metadata.id.as_user()?,
 554                        title: prompt_metadata.title.map(|title| title.to_string()),
 555                        contents,
 556                    }),
 557                    Err(_err) => {
 558                        // TODO: show error message
 559                        // this.update(cx, |_, cx| {
 560                        //     cx.emit(RulesLoadingError {
 561                        //         message: format!("{err:?}").into(),
 562                        //     });
 563                        // })
 564                        // .ok();
 565                        None
 566                    }
 567                })
 568                .collect::<Vec<_>>();
 569
 570            ProjectContext::new(worktrees, default_user_rules)
 571        })
 572    }
 573
 574    fn load_worktree_info_for_system_prompt(
 575        worktree: Entity<Worktree>,
 576        project: Entity<Project>,
 577        cx: &mut App,
 578    ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
 579        let tree = worktree.read(cx);
 580        let root_name = tree.root_name_str().into();
 581        let abs_path = tree.abs_path();
 582
 583        let mut context = WorktreeContext {
 584            root_name,
 585            abs_path,
 586            rules_file: None,
 587        };
 588
 589        let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
 590        let Some(rules_task) = rules_task else {
 591            return Task::ready((context, None));
 592        };
 593
 594        cx.spawn(async move |_| {
 595            let (rules_file, rules_file_error) = match rules_task.await {
 596                Ok(rules_file) => (Some(rules_file), None),
 597                Err(err) => (
 598                    None,
 599                    Some(RulesLoadingError {
 600                        message: format!("{err}").into(),
 601                    }),
 602                ),
 603            };
 604            context.rules_file = rules_file;
 605            (context, rules_file_error)
 606        })
 607    }
 608
 609    fn load_worktree_rules_file(
 610        worktree: Entity<Worktree>,
 611        project: Entity<Project>,
 612        cx: &mut App,
 613    ) -> Option<Task<Result<RulesFileContext>>> {
 614        let worktree = worktree.read(cx);
 615        let worktree_id = worktree.id();
 616        let selected_rules_file = RULES_FILE_NAMES
 617            .into_iter()
 618            .filter_map(|name| {
 619                worktree
 620                    .entry_for_path(RelPath::unix(name).unwrap())
 621                    .filter(|entry| entry.is_file())
 622                    .map(|entry| entry.path.clone())
 623            })
 624            .next();
 625
 626        // Note that Cline supports `.clinerules` being a directory, but that is not currently
 627        // supported. This doesn't seem to occur often in GitHub repositories.
 628        selected_rules_file.map(|path_in_worktree| {
 629            let project_path = ProjectPath {
 630                worktree_id,
 631                path: path_in_worktree.clone(),
 632            };
 633            let buffer_task =
 634                project.update(cx, |project, cx| project.open_buffer(project_path, cx));
 635            let rope_task = cx.spawn(async move |cx| {
 636                let buffer = buffer_task.await?;
 637                let (project_entry_id, rope) = buffer.read_with(cx, |buffer, cx| {
 638                    let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
 639                    anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
 640                })?;
 641                anyhow::Ok((project_entry_id, rope))
 642            });
 643            // Build a string from the rope on a background thread.
 644            cx.background_spawn(async move {
 645                let (project_entry_id, rope) = rope_task.await?;
 646                anyhow::Ok(RulesFileContext {
 647                    path_in_worktree,
 648                    text: rope.to_string().trim().to_string(),
 649                    project_entry_id: project_entry_id.to_usize(),
 650                })
 651            })
 652        })
 653    }
 654
 655    fn handle_thread_title_updated(
 656        &mut self,
 657        thread: Entity<Thread>,
 658        _: &TitleUpdated,
 659        cx: &mut Context<Self>,
 660    ) {
 661        let session_id = thread.read(cx).id();
 662        let Some(session) = self.sessions.get(session_id) else {
 663            return;
 664        };
 665        let thread = thread.downgrade();
 666        let acp_thread = session.acp_thread.downgrade();
 667        cx.spawn(async move |_, cx| {
 668            let title = thread.read_with(cx, |thread, _| thread.title())?;
 669            let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
 670            task.await
 671        })
 672        .detach_and_log_err(cx);
 673    }
 674
 675    fn handle_thread_token_usage_updated(
 676        &mut self,
 677        thread: Entity<Thread>,
 678        usage: &TokenUsageUpdated,
 679        cx: &mut Context<Self>,
 680    ) {
 681        let Some(session) = self.sessions.get(thread.read(cx).id()) else {
 682            return;
 683        };
 684        session.acp_thread.update(cx, |acp_thread, cx| {
 685            acp_thread.update_token_usage(usage.0.clone(), cx);
 686        });
 687    }
 688
 689    fn handle_project_event(
 690        &mut self,
 691        project: Entity<Project>,
 692        event: &project::Event,
 693        _cx: &mut Context<Self>,
 694    ) {
 695        let project_id = project.entity_id();
 696        let Some(state) = self.projects.get_mut(&project_id) else {
 697            return;
 698        };
 699        match event {
 700            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
 701                state.project_context_needs_refresh.send(()).ok();
 702            }
 703            project::Event::WorktreeUpdatedEntries(_, items) => {
 704                if items.iter().any(|(path, _, _)| {
 705                    RULES_FILE_NAMES
 706                        .iter()
 707                        .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
 708                }) {
 709                    state.project_context_needs_refresh.send(()).ok();
 710                }
 711            }
 712            _ => {}
 713        }
 714    }
 715
 716    fn handle_prompts_updated_event(
 717        &mut self,
 718        _prompt_store: Entity<PromptStore>,
 719        _event: &prompt_store::PromptsUpdatedEvent,
 720        _cx: &mut Context<Self>,
 721    ) {
 722        for state in self.projects.values_mut() {
 723            state.project_context_needs_refresh.send(()).ok();
 724        }
 725    }
 726
 727    fn handle_models_updated_event(
 728        &mut self,
 729        _registry: Entity<LanguageModelRegistry>,
 730        _event: &language_model::Event,
 731        cx: &mut Context<Self>,
 732    ) {
 733        self.models.refresh_list(cx);
 734
 735        let registry = LanguageModelRegistry::read_global(cx);
 736        let default_model = registry.default_model().map(|m| m.model);
 737        let summarization_model = registry.thread_summary_model().map(|m| m.model);
 738
 739        for session in self.sessions.values_mut() {
 740            session.thread.update(cx, |thread, cx| {
 741                if thread.model().is_none()
 742                    && let Some(model) = default_model.clone()
 743                {
 744                    thread.set_model(model, cx);
 745                    cx.notify();
 746                }
 747                thread.set_summarization_model(summarization_model.clone(), cx);
 748            });
 749        }
 750    }
 751
 752    fn handle_context_server_store_updated(
 753        &mut self,
 754        store: Entity<project::context_server_store::ContextServerStore>,
 755        _event: &project::context_server_store::ServerStatusChangedEvent,
 756        cx: &mut Context<Self>,
 757    ) {
 758        let project_id = self.projects.iter().find_map(|(id, state)| {
 759            if *state.context_server_registry.read(cx).server_store() == store {
 760                Some(*id)
 761            } else {
 762                None
 763            }
 764        });
 765        if let Some(project_id) = project_id {
 766            self.update_available_commands_for_project(project_id, cx);
 767        }
 768    }
 769
 770    fn handle_context_server_registry_event(
 771        &mut self,
 772        registry: Entity<ContextServerRegistry>,
 773        event: &ContextServerRegistryEvent,
 774        cx: &mut Context<Self>,
 775    ) {
 776        match event {
 777            ContextServerRegistryEvent::ToolsChanged => {}
 778            ContextServerRegistryEvent::PromptsChanged => {
 779                let project_id = self.projects.iter().find_map(|(id, state)| {
 780                    if state.context_server_registry == registry {
 781                        Some(*id)
 782                    } else {
 783                        None
 784                    }
 785                });
 786                if let Some(project_id) = project_id {
 787                    self.update_available_commands_for_project(project_id, cx);
 788                }
 789            }
 790        }
 791    }
 792
 793    fn update_available_commands_for_project(&self, project_id: EntityId, cx: &mut Context<Self>) {
 794        let available_commands =
 795            Self::build_available_commands_for_project(self.projects.get(&project_id), cx);
 796        for session in self.sessions.values() {
 797            if session.project_id != project_id {
 798                continue;
 799            }
 800            session.acp_thread.update(cx, |thread, cx| {
 801                thread
 802                    .handle_session_update(
 803                        acp::SessionUpdate::AvailableCommandsUpdate(
 804                            acp::AvailableCommandsUpdate::new(available_commands.clone()),
 805                        ),
 806                        cx,
 807                    )
 808                    .log_err();
 809            });
 810        }
 811    }
 812
 813    fn build_available_commands_for_project(
 814        project_state: Option<&ProjectState>,
 815        cx: &App,
 816    ) -> Vec<acp::AvailableCommand> {
 817        let Some(state) = project_state else {
 818            return vec![];
 819        };
 820        let registry = state.context_server_registry.read(cx);
 821
 822        let mut prompt_name_counts: HashMap<&str, usize> = HashMap::default();
 823        for context_server_prompt in registry.prompts() {
 824            *prompt_name_counts
 825                .entry(context_server_prompt.prompt.name.as_str())
 826                .or_insert(0) += 1;
 827        }
 828
 829        registry
 830            .prompts()
 831            .flat_map(|context_server_prompt| {
 832                let prompt = &context_server_prompt.prompt;
 833
 834                let should_prefix = prompt_name_counts
 835                    .get(prompt.name.as_str())
 836                    .copied()
 837                    .unwrap_or(0)
 838                    > 1;
 839
 840                let name = if should_prefix {
 841                    format!("{}.{}", context_server_prompt.server_id, prompt.name)
 842                } else {
 843                    prompt.name.clone()
 844                };
 845
 846                let mut command = acp::AvailableCommand::new(
 847                    name,
 848                    prompt.description.clone().unwrap_or_default(),
 849                );
 850
 851                match prompt.arguments.as_deref() {
 852                    Some([arg]) => {
 853                        let hint = format!("<{}>", arg.name);
 854
 855                        command = command.input(acp::AvailableCommandInput::Unstructured(
 856                            acp::UnstructuredCommandInput::new(hint),
 857                        ));
 858                    }
 859                    Some([]) | None => {}
 860                    Some(_) => {
 861                        // skip >1 argument commands since we don't support them yet
 862                        return None;
 863                    }
 864                }
 865
 866                Some(command)
 867            })
 868            .collect()
 869    }
 870
 871    pub fn load_thread(
 872        &mut self,
 873        id: acp::SessionId,
 874        project: Entity<Project>,
 875        cx: &mut Context<Self>,
 876    ) -> Task<Result<Entity<Thread>>> {
 877        let database_future = ThreadsDatabase::connect(cx);
 878        cx.spawn(async move |this, cx| {
 879            let database = database_future.await.map_err(|err| anyhow!(err))?;
 880            let db_thread = database
 881                .load_thread(id.clone())
 882                .await?
 883                .with_context(|| format!("no thread found with ID: {id:?}"))?;
 884
 885            this.update(cx, |this, cx| {
 886                let project_id = this.get_or_create_project_state(&project, cx);
 887                let project_state = this
 888                    .projects
 889                    .get(&project_id)
 890                    .context("project state not found")?;
 891                let summarization_model = LanguageModelRegistry::read_global(cx)
 892                    .thread_summary_model()
 893                    .map(|c| c.model);
 894
 895                Ok(cx.new(|cx| {
 896                    let mut thread = Thread::from_db(
 897                        id.clone(),
 898                        db_thread,
 899                        project_state.project.clone(),
 900                        project_state.project_context.clone(),
 901                        project_state.context_server_registry.clone(),
 902                        this.templates.clone(),
 903                        cx,
 904                    );
 905                    thread.set_summarization_model(summarization_model, cx);
 906                    thread
 907                }))
 908            })?
 909        })
 910    }
 911
 912    pub fn open_thread(
 913        &mut self,
 914        id: acp::SessionId,
 915        project: Entity<Project>,
 916        cx: &mut Context<Self>,
 917    ) -> Task<Result<Entity<AcpThread>>> {
 918        if let Some(session) = self.sessions.get(&id) {
 919            return Task::ready(Ok(session.acp_thread.clone()));
 920        }
 921
 922        let task = self.load_thread(id, project.clone(), cx);
 923        cx.spawn(async move |this, cx| {
 924            let thread = task.await?;
 925            let acp_thread = this.update(cx, |this, cx| {
 926                let project_id = this.get_or_create_project_state(&project, cx);
 927                this.register_session(thread.clone(), project_id, cx)
 928            })?;
 929            let events = thread.update(cx, |thread, cx| thread.replay(cx));
 930            cx.update(|cx| {
 931                NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
 932            })
 933            .await?;
 934            Ok(acp_thread)
 935        })
 936    }
 937
 938    pub fn thread_summary(
 939        &mut self,
 940        id: acp::SessionId,
 941        project: Entity<Project>,
 942        cx: &mut Context<Self>,
 943    ) -> Task<Result<SharedString>> {
 944        let thread = self.open_thread(id.clone(), project, cx);
 945        cx.spawn(async move |this, cx| {
 946            let acp_thread = thread.await?;
 947            let result = this
 948                .update(cx, |this, cx| {
 949                    this.sessions
 950                        .get(&id)
 951                        .unwrap()
 952                        .thread
 953                        .update(cx, |thread, cx| thread.summary(cx))
 954                })?
 955                .await
 956                .context("Failed to generate summary")?;
 957            drop(acp_thread);
 958            Ok(result)
 959        })
 960    }
 961
 962    fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
 963        if thread.read(cx).is_empty() {
 964            return;
 965        }
 966
 967        let id = thread.read(cx).id().clone();
 968        let Some(session) = self.sessions.get_mut(&id) else {
 969            return;
 970        };
 971
 972        let project_id = session.project_id;
 973        let Some(state) = self.projects.get(&project_id) else {
 974            return;
 975        };
 976
 977        let folder_paths = PathList::new(
 978            &state
 979                .project
 980                .read(cx)
 981                .visible_worktrees(cx)
 982                .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
 983                .collect::<Vec<_>>(),
 984        );
 985
 986        let draft_prompt = session.acp_thread.read(cx).draft_prompt().map(Vec::from);
 987        let database_future = ThreadsDatabase::connect(cx);
 988        let db_thread = thread.update(cx, |thread, cx| {
 989            thread.set_draft_prompt(draft_prompt);
 990            thread.to_db(cx)
 991        });
 992        let thread_store = self.thread_store.clone();
 993        session.pending_save = cx.spawn(async move |_, cx| {
 994            let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
 995                return;
 996            };
 997            let db_thread = db_thread.await;
 998            database
 999                .save_thread(id, db_thread, folder_paths)
1000                .await
1001                .log_err();
1002            thread_store.update(cx, |store, cx| store.reload(cx));
1003        });
1004    }
1005
1006    fn send_mcp_prompt(
1007        &self,
1008        message_id: UserMessageId,
1009        session_id: acp::SessionId,
1010        prompt_name: String,
1011        server_id: ContextServerId,
1012        arguments: HashMap<String, String>,
1013        original_content: Vec<acp::ContentBlock>,
1014        cx: &mut Context<Self>,
1015    ) -> Task<Result<acp::PromptResponse>> {
1016        let Some(state) = self.session_project_state(&session_id) else {
1017            return Task::ready(Err(anyhow!("Project state not found for session")));
1018        };
1019        let server_store = state
1020            .context_server_registry
1021            .read(cx)
1022            .server_store()
1023            .clone();
1024        let path_style = state.project.read(cx).path_style(cx);
1025
1026        cx.spawn(async move |this, cx| {
1027            let prompt =
1028                crate::get_prompt(&server_store, &server_id, &prompt_name, arguments, cx).await?;
1029
1030            let (acp_thread, thread) = this.update(cx, |this, _cx| {
1031                let session = this
1032                    .sessions
1033                    .get(&session_id)
1034                    .context("Failed to get session")?;
1035                anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
1036            })??;
1037
1038            let mut last_is_user = true;
1039
1040            thread.update(cx, |thread, cx| {
1041                thread.push_acp_user_block(
1042                    message_id,
1043                    original_content.into_iter().skip(1),
1044                    path_style,
1045                    cx,
1046                );
1047            });
1048
1049            for message in prompt.messages {
1050                let context_server::types::PromptMessage { role, content } = message;
1051                let block = mcp_message_content_to_acp_content_block(content);
1052
1053                match role {
1054                    context_server::types::Role::User => {
1055                        let id = acp_thread::UserMessageId::new();
1056
1057                        acp_thread.update(cx, |acp_thread, cx| {
1058                            acp_thread.push_user_content_block_with_indent(
1059                                Some(id.clone()),
1060                                block.clone(),
1061                                true,
1062                                cx,
1063                            );
1064                        });
1065
1066                        thread.update(cx, |thread, cx| {
1067                            thread.push_acp_user_block(id, [block], path_style, cx);
1068                        });
1069                    }
1070                    context_server::types::Role::Assistant => {
1071                        acp_thread.update(cx, |acp_thread, cx| {
1072                            acp_thread.push_assistant_content_block_with_indent(
1073                                block.clone(),
1074                                false,
1075                                true,
1076                                cx,
1077                            );
1078                        });
1079
1080                        thread.update(cx, |thread, cx| {
1081                            thread.push_acp_agent_block(block, cx);
1082                        });
1083                    }
1084                }
1085
1086                last_is_user = role == context_server::types::Role::User;
1087            }
1088
1089            let response_stream = thread.update(cx, |thread, cx| {
1090                if last_is_user {
1091                    thread.send_existing(cx)
1092                } else {
1093                    // Resume if MCP prompt did not end with a user message
1094                    thread.resume(cx)
1095                }
1096            })?;
1097
1098            cx.update(|cx| {
1099                NativeAgentConnection::handle_thread_events(
1100                    response_stream,
1101                    acp_thread.downgrade(),
1102                    cx,
1103                )
1104            })
1105            .await
1106        })
1107    }
1108}
1109
1110/// Wrapper struct that implements the AgentConnection trait
1111#[derive(Clone)]
1112pub struct NativeAgentConnection(pub Entity<NativeAgent>);
1113
1114impl NativeAgentConnection {
1115    pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
1116        self.0
1117            .read(cx)
1118            .sessions
1119            .get(session_id)
1120            .map(|session| session.thread.clone())
1121    }
1122
1123    pub fn load_thread(
1124        &self,
1125        id: acp::SessionId,
1126        project: Entity<Project>,
1127        cx: &mut App,
1128    ) -> Task<Result<Entity<Thread>>> {
1129        self.0
1130            .update(cx, |this, cx| this.load_thread(id, project, cx))
1131    }
1132
1133    fn run_turn(
1134        &self,
1135        session_id: acp::SessionId,
1136        cx: &mut App,
1137        f: impl 'static
1138        + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
1139    ) -> Task<Result<acp::PromptResponse>> {
1140        let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
1141            agent
1142                .sessions
1143                .get_mut(&session_id)
1144                .map(|s| (s.thread.clone(), s.acp_thread.clone()))
1145        }) else {
1146            return Task::ready(Err(anyhow!("Session not found")));
1147        };
1148        log::debug!("Found session for: {}", session_id);
1149
1150        let response_stream = match f(thread, cx) {
1151            Ok(stream) => stream,
1152            Err(err) => return Task::ready(Err(err)),
1153        };
1154        Self::handle_thread_events(response_stream, acp_thread.downgrade(), cx)
1155    }
1156
1157    fn handle_thread_events(
1158        mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
1159        acp_thread: WeakEntity<AcpThread>,
1160        cx: &App,
1161    ) -> Task<Result<acp::PromptResponse>> {
1162        cx.spawn(async move |cx| {
1163            // Handle response stream and forward to session.acp_thread
1164            while let Some(result) = events.next().await {
1165                match result {
1166                    Ok(event) => {
1167                        log::trace!("Received completion event: {:?}", event);
1168
1169                        match event {
1170                            ThreadEvent::UserMessage(message) => {
1171                                acp_thread.update(cx, |thread, cx| {
1172                                    for content in message.content {
1173                                        thread.push_user_content_block(
1174                                            Some(message.id.clone()),
1175                                            content.into(),
1176                                            cx,
1177                                        );
1178                                    }
1179                                })?;
1180                            }
1181                            ThreadEvent::AgentText(text) => {
1182                                acp_thread.update(cx, |thread, cx| {
1183                                    thread.push_assistant_content_block(text.into(), false, cx)
1184                                })?;
1185                            }
1186                            ThreadEvent::AgentThinking(text) => {
1187                                acp_thread.update(cx, |thread, cx| {
1188                                    thread.push_assistant_content_block(text.into(), true, cx)
1189                                })?;
1190                            }
1191                            ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
1192                                tool_call,
1193                                options,
1194                                response,
1195                                context: _,
1196                            }) => {
1197                                let outcome_task = acp_thread.update(cx, |thread, cx| {
1198                                    thread.request_tool_call_authorization(tool_call, options, cx)
1199                                })??;
1200                                cx.background_spawn(async move {
1201                                    if let acp_thread::RequestPermissionOutcome::Selected(outcome) =
1202                                        outcome_task.await
1203                                    {
1204                                        response
1205                                            .send(outcome)
1206                                            .map(|_| anyhow!("authorization receiver was dropped"))
1207                                            .log_err();
1208                                    }
1209                                })
1210                                .detach();
1211                            }
1212                            ThreadEvent::ToolCall(tool_call) => {
1213                                acp_thread.update(cx, |thread, cx| {
1214                                    thread.upsert_tool_call(tool_call, cx)
1215                                })??;
1216                            }
1217                            ThreadEvent::ToolCallUpdate(update) => {
1218                                acp_thread.update(cx, |thread, cx| {
1219                                    thread.update_tool_call(update, cx)
1220                                })??;
1221                            }
1222                            ThreadEvent::SubagentSpawned(session_id) => {
1223                                acp_thread.update(cx, |thread, cx| {
1224                                    thread.subagent_spawned(session_id, cx);
1225                                })?;
1226                            }
1227                            ThreadEvent::Retry(status) => {
1228                                acp_thread.update(cx, |thread, cx| {
1229                                    thread.update_retry_status(status, cx)
1230                                })?;
1231                            }
1232                            ThreadEvent::Stop(stop_reason) => {
1233                                log::debug!("Assistant message complete: {:?}", stop_reason);
1234                                return Ok(acp::PromptResponse::new(stop_reason));
1235                            }
1236                        }
1237                    }
1238                    Err(e) => {
1239                        log::error!("Error in model response stream: {:?}", e);
1240                        return Err(e);
1241                    }
1242                }
1243            }
1244
1245            log::debug!("Response stream completed");
1246            anyhow::Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
1247        })
1248    }
1249}
1250
1251struct Command<'a> {
1252    prompt_name: &'a str,
1253    arg_value: &'a str,
1254    explicit_server_id: Option<&'a str>,
1255}
1256
1257impl<'a> Command<'a> {
1258    fn parse(prompt: &'a [acp::ContentBlock]) -> Option<Self> {
1259        let acp::ContentBlock::Text(text_content) = prompt.first()? else {
1260            return None;
1261        };
1262        let text = text_content.text.trim();
1263        let command = text.strip_prefix('/')?;
1264        let (command, arg_value) = command
1265            .split_once(char::is_whitespace)
1266            .unwrap_or((command, ""));
1267
1268        if let Some((server_id, prompt_name)) = command.split_once('.') {
1269            Some(Self {
1270                prompt_name,
1271                arg_value,
1272                explicit_server_id: Some(server_id),
1273            })
1274        } else {
1275            Some(Self {
1276                prompt_name: command,
1277                arg_value,
1278                explicit_server_id: None,
1279            })
1280        }
1281    }
1282}
1283
1284struct NativeAgentModelSelector {
1285    session_id: acp::SessionId,
1286    connection: NativeAgentConnection,
1287}
1288
1289impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
1290    fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
1291        log::debug!("NativeAgentConnection::list_models called");
1292        let list = self.connection.0.read(cx).models.model_list.clone();
1293        Task::ready(if list.is_empty() {
1294            Err(anyhow::anyhow!("No models available"))
1295        } else {
1296            Ok(list)
1297        })
1298    }
1299
1300    fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
1301        log::debug!(
1302            "Setting model for session {}: {}",
1303            self.session_id,
1304            model_id
1305        );
1306        let Some(thread) = self
1307            .connection
1308            .0
1309            .read(cx)
1310            .sessions
1311            .get(&self.session_id)
1312            .map(|session| session.thread.clone())
1313        else {
1314            return Task::ready(Err(anyhow!("Session not found")));
1315        };
1316
1317        let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
1318            return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
1319        };
1320
1321        // We want to reset the effort level when switching models, as the currently-selected effort level may
1322        // not be compatible.
1323        let effort = model
1324            .default_effort_level()
1325            .map(|effort_level| effort_level.value.to_string());
1326
1327        thread.update(cx, |thread, cx| {
1328            thread.set_model(model.clone(), cx);
1329            thread.set_thinking_effort(effort.clone(), cx);
1330            thread.set_thinking_enabled(model.supports_thinking(), cx);
1331        });
1332
1333        update_settings_file(
1334            self.connection.0.read(cx).fs.clone(),
1335            cx,
1336            move |settings, cx| {
1337                let provider = model.provider_id().0.to_string();
1338                let model = model.id().0.to_string();
1339                let enable_thinking = thread.read(cx).thinking_enabled();
1340                settings
1341                    .agent
1342                    .get_or_insert_default()
1343                    .set_model(LanguageModelSelection {
1344                        provider: provider.into(),
1345                        model,
1346                        enable_thinking,
1347                        effort,
1348                    });
1349            },
1350        );
1351
1352        Task::ready(Ok(()))
1353    }
1354
1355    fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
1356        let Some(thread) = self
1357            .connection
1358            .0
1359            .read(cx)
1360            .sessions
1361            .get(&self.session_id)
1362            .map(|session| session.thread.clone())
1363        else {
1364            return Task::ready(Err(anyhow!("Session not found")));
1365        };
1366        let Some(model) = thread.read(cx).model() else {
1367            return Task::ready(Err(anyhow!("Model not found")));
1368        };
1369        let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
1370        else {
1371            return Task::ready(Err(anyhow!("Provider not found")));
1372        };
1373        Task::ready(Ok(LanguageModels::map_language_model_to_info(
1374            model, &provider,
1375        )))
1376    }
1377
1378    fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
1379        Some(self.connection.0.read(cx).models.watch())
1380    }
1381
1382    fn should_render_footer(&self) -> bool {
1383        true
1384    }
1385}
1386
1387pub static ZED_AGENT_ID: LazyLock<AgentId> = LazyLock::new(|| AgentId::new("Zed Agent"));
1388
1389impl acp_thread::AgentConnection for NativeAgentConnection {
1390    fn agent_id(&self) -> AgentId {
1391        ZED_AGENT_ID.clone()
1392    }
1393
1394    fn telemetry_id(&self) -> SharedString {
1395        "zed".into()
1396    }
1397
1398    fn new_session(
1399        self: Rc<Self>,
1400        project: Entity<Project>,
1401        work_dirs: PathList,
1402        cx: &mut App,
1403    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1404        log::debug!("Creating new thread for project at: {work_dirs:?}");
1405        Task::ready(Ok(self
1406            .0
1407            .update(cx, |agent, cx| agent.new_session(project, cx))))
1408    }
1409
1410    fn supports_load_session(&self) -> bool {
1411        true
1412    }
1413
1414    fn load_session(
1415        self: Rc<Self>,
1416        session_id: acp::SessionId,
1417        project: Entity<Project>,
1418        _work_dirs: PathList,
1419        _title: Option<SharedString>,
1420        cx: &mut App,
1421    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
1422        self.0
1423            .update(cx, |agent, cx| agent.open_thread(session_id, project, cx))
1424    }
1425
1426    fn supports_close_session(&self) -> bool {
1427        true
1428    }
1429
1430    fn close_session(
1431        self: Rc<Self>,
1432        session_id: &acp::SessionId,
1433        cx: &mut App,
1434    ) -> Task<Result<()>> {
1435        self.0.update(cx, |agent, cx| {
1436            let Some(session) = agent.sessions.remove(session_id) else {
1437                return;
1438            };
1439            let project_id = session.project_id;
1440            agent.save_thread(session.thread, cx);
1441
1442            let has_remaining = agent.sessions.values().any(|s| s.project_id == project_id);
1443            if !has_remaining {
1444                agent.projects.remove(&project_id);
1445            }
1446        });
1447        Task::ready(Ok(()))
1448    }
1449
1450    fn auth_methods(&self) -> &[acp::AuthMethod] {
1451        &[] // No auth for in-process
1452    }
1453
1454    fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1455        Task::ready(Ok(()))
1456    }
1457
1458    fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1459        Some(Rc::new(NativeAgentModelSelector {
1460            session_id: session_id.clone(),
1461            connection: self.clone(),
1462        }) as Rc<dyn AgentModelSelector>)
1463    }
1464
1465    fn prompt(
1466        &self,
1467        id: Option<acp_thread::UserMessageId>,
1468        params: acp::PromptRequest,
1469        cx: &mut App,
1470    ) -> Task<Result<acp::PromptResponse>> {
1471        let id = id.expect("UserMessageId is required");
1472        let session_id = params.session_id.clone();
1473        log::info!("Received prompt request for session: {}", session_id);
1474        log::debug!("Prompt blocks count: {}", params.prompt.len());
1475
1476        let Some(project_state) = self.0.read(cx).session_project_state(&session_id) else {
1477            return Task::ready(Err(anyhow::anyhow!("Session not found")));
1478        };
1479
1480        if let Some(parsed_command) = Command::parse(&params.prompt) {
1481            let registry = project_state.context_server_registry.read(cx);
1482
1483            let explicit_server_id = parsed_command
1484                .explicit_server_id
1485                .map(|server_id| ContextServerId(server_id.into()));
1486
1487            if let Some(prompt) =
1488                registry.find_prompt(explicit_server_id.as_ref(), parsed_command.prompt_name)
1489            {
1490                let arguments = if !parsed_command.arg_value.is_empty()
1491                    && let Some(arg_name) = prompt
1492                        .prompt
1493                        .arguments
1494                        .as_ref()
1495                        .and_then(|args| args.first())
1496                        .map(|arg| arg.name.clone())
1497                {
1498                    HashMap::from_iter([(arg_name, parsed_command.arg_value.to_string())])
1499                } else {
1500                    Default::default()
1501                };
1502
1503                let prompt_name = prompt.prompt.name.clone();
1504                let server_id = prompt.server_id.clone();
1505
1506                return self.0.update(cx, |agent, cx| {
1507                    agent.send_mcp_prompt(
1508                        id,
1509                        session_id.clone(),
1510                        prompt_name,
1511                        server_id,
1512                        arguments,
1513                        params.prompt,
1514                        cx,
1515                    )
1516                });
1517            }
1518        };
1519
1520        let path_style = project_state.project.read(cx).path_style(cx);
1521
1522        self.run_turn(session_id, cx, move |thread, cx| {
1523            let content: Vec<UserMessageContent> = params
1524                .prompt
1525                .into_iter()
1526                .map(|block| UserMessageContent::from_content_block(block, path_style))
1527                .collect::<Vec<_>>();
1528            log::debug!("Converted prompt to message: {} chars", content.len());
1529            log::debug!("Message id: {:?}", id);
1530            log::debug!("Message content: {:?}", content);
1531
1532            thread.update(cx, |thread, cx| thread.send(id, content, cx))
1533        })
1534    }
1535
1536    fn retry(
1537        &self,
1538        session_id: &acp::SessionId,
1539        _cx: &App,
1540    ) -> Option<Rc<dyn acp_thread::AgentSessionRetry>> {
1541        Some(Rc::new(NativeAgentSessionRetry {
1542            connection: self.clone(),
1543            session_id: session_id.clone(),
1544        }) as _)
1545    }
1546
1547    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1548        log::info!("Cancelling on session: {}", session_id);
1549        self.0.update(cx, |agent, cx| {
1550            if let Some(session) = agent.sessions.get(session_id) {
1551                session
1552                    .thread
1553                    .update(cx, |thread, cx| thread.cancel(cx))
1554                    .detach();
1555            }
1556        });
1557    }
1558
1559    fn truncate(
1560        &self,
1561        session_id: &acp::SessionId,
1562        cx: &App,
1563    ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1564        self.0.read_with(cx, |agent, _cx| {
1565            agent.sessions.get(session_id).map(|session| {
1566                Rc::new(NativeAgentSessionTruncate {
1567                    thread: session.thread.clone(),
1568                    acp_thread: session.acp_thread.downgrade(),
1569                }) as _
1570            })
1571        })
1572    }
1573
1574    fn set_title(
1575        &self,
1576        session_id: &acp::SessionId,
1577        cx: &App,
1578    ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1579        self.0.read_with(cx, |agent, _cx| {
1580            agent
1581                .sessions
1582                .get(session_id)
1583                .filter(|s| !s.thread.read(cx).is_subagent())
1584                .map(|session| {
1585                    Rc::new(NativeAgentSessionSetTitle {
1586                        thread: session.thread.clone(),
1587                    }) as _
1588                })
1589        })
1590    }
1591
1592    fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
1593        let thread_store = self.0.read(cx).thread_store.clone();
1594        Some(Rc::new(NativeAgentSessionList::new(thread_store, cx)) as _)
1595    }
1596
1597    fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1598        Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1599    }
1600
1601    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1602        self
1603    }
1604}
1605
1606impl acp_thread::AgentTelemetry for NativeAgentConnection {
1607    fn thread_data(
1608        &self,
1609        session_id: &acp::SessionId,
1610        cx: &mut App,
1611    ) -> Task<Result<serde_json::Value>> {
1612        let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1613            return Task::ready(Err(anyhow!("Session not found")));
1614        };
1615
1616        let task = session.thread.read(cx).to_db(cx);
1617        cx.background_spawn(async move {
1618            serde_json::to_value(task.await).context("Failed to serialize thread")
1619        })
1620    }
1621}
1622
1623pub struct NativeAgentSessionList {
1624    thread_store: Entity<ThreadStore>,
1625    updates_tx: smol::channel::Sender<acp_thread::SessionListUpdate>,
1626    updates_rx: smol::channel::Receiver<acp_thread::SessionListUpdate>,
1627    _subscription: Subscription,
1628}
1629
1630impl NativeAgentSessionList {
1631    fn new(thread_store: Entity<ThreadStore>, cx: &mut App) -> Self {
1632        let (tx, rx) = smol::channel::unbounded();
1633        let this_tx = tx.clone();
1634        let subscription = cx.observe(&thread_store, move |_, _| {
1635            this_tx
1636                .try_send(acp_thread::SessionListUpdate::Refresh)
1637                .ok();
1638        });
1639        Self {
1640            thread_store,
1641            updates_tx: tx,
1642            updates_rx: rx,
1643            _subscription: subscription,
1644        }
1645    }
1646
1647    pub fn thread_store(&self) -> &Entity<ThreadStore> {
1648        &self.thread_store
1649    }
1650}
1651
1652impl AgentSessionList for NativeAgentSessionList {
1653    fn list_sessions(
1654        &self,
1655        _request: AgentSessionListRequest,
1656        cx: &mut App,
1657    ) -> Task<Result<AgentSessionListResponse>> {
1658        let sessions = self
1659            .thread_store
1660            .read(cx)
1661            .entries()
1662            .map(|entry| AgentSessionInfo::from(&entry))
1663            .collect();
1664        Task::ready(Ok(AgentSessionListResponse::new(sessions)))
1665    }
1666
1667    fn supports_delete(&self) -> bool {
1668        true
1669    }
1670
1671    fn delete_session(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<()>> {
1672        self.thread_store
1673            .update(cx, |store, cx| store.delete_thread(session_id.clone(), cx))
1674    }
1675
1676    fn delete_sessions(&self, cx: &mut App) -> Task<Result<()>> {
1677        self.thread_store
1678            .update(cx, |store, cx| store.delete_threads(cx))
1679    }
1680
1681    fn watch(
1682        &self,
1683        _cx: &mut App,
1684    ) -> Option<smol::channel::Receiver<acp_thread::SessionListUpdate>> {
1685        Some(self.updates_rx.clone())
1686    }
1687
1688    fn notify_refresh(&self) {
1689        self.updates_tx
1690            .try_send(acp_thread::SessionListUpdate::Refresh)
1691            .ok();
1692    }
1693
1694    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1695        self
1696    }
1697}
1698
1699struct NativeAgentSessionTruncate {
1700    thread: Entity<Thread>,
1701    acp_thread: WeakEntity<AcpThread>,
1702}
1703
1704impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1705    fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1706        match self.thread.update(cx, |thread, cx| {
1707            thread.truncate(message_id.clone(), cx)?;
1708            Ok(thread.latest_token_usage())
1709        }) {
1710            Ok(usage) => {
1711                self.acp_thread
1712                    .update(cx, |thread, cx| {
1713                        thread.update_token_usage(usage, cx);
1714                    })
1715                    .ok();
1716                Task::ready(Ok(()))
1717            }
1718            Err(error) => Task::ready(Err(error)),
1719        }
1720    }
1721}
1722
1723struct NativeAgentSessionRetry {
1724    connection: NativeAgentConnection,
1725    session_id: acp::SessionId,
1726}
1727
1728impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
1729    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1730        self.connection
1731            .run_turn(self.session_id.clone(), cx, |thread, cx| {
1732                thread.update(cx, |thread, cx| thread.resume(cx))
1733            })
1734    }
1735}
1736
1737struct NativeAgentSessionSetTitle {
1738    thread: Entity<Thread>,
1739}
1740
1741impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1742    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1743        self.thread
1744            .update(cx, |thread, cx| thread.set_title(title, cx));
1745        Task::ready(Ok(()))
1746    }
1747}
1748
1749pub struct NativeThreadEnvironment {
1750    agent: WeakEntity<NativeAgent>,
1751    thread: WeakEntity<Thread>,
1752    acp_thread: WeakEntity<AcpThread>,
1753}
1754
1755impl NativeThreadEnvironment {
1756    pub(crate) fn create_subagent_thread(
1757        &self,
1758        label: String,
1759        cx: &mut App,
1760    ) -> Result<Rc<dyn SubagentHandle>> {
1761        let Some(parent_thread_entity) = self.thread.upgrade() else {
1762            anyhow::bail!("Parent thread no longer exists".to_string());
1763        };
1764        let parent_thread = parent_thread_entity.read(cx);
1765        let current_depth = parent_thread.depth();
1766        let parent_session_id = parent_thread.id().clone();
1767
1768        if current_depth >= MAX_SUBAGENT_DEPTH {
1769            return Err(anyhow!(
1770                "Maximum subagent depth ({}) reached",
1771                MAX_SUBAGENT_DEPTH
1772            ));
1773        }
1774
1775        let subagent_thread: Entity<Thread> = cx.new(|cx| {
1776            let mut thread = Thread::new_subagent(&parent_thread_entity, cx);
1777            thread.set_title(label.into(), cx);
1778            thread
1779        });
1780
1781        let session_id = subagent_thread.read(cx).id().clone();
1782
1783        let acp_thread = self
1784            .agent
1785            .update(cx, |agent, cx| -> Result<Entity<AcpThread>> {
1786                let project_id = agent
1787                    .sessions
1788                    .get(&parent_session_id)
1789                    .map(|s| s.project_id)
1790                    .context("parent session not found")?;
1791                Ok(agent.register_session(subagent_thread.clone(), project_id, cx))
1792            })??;
1793
1794        let depth = current_depth + 1;
1795
1796        telemetry::event!(
1797            "Subagent Started",
1798            session = parent_thread_entity.read(cx).id().to_string(),
1799            subagent_session = session_id.to_string(),
1800            depth,
1801            is_resumed = false,
1802        );
1803
1804        self.prompt_subagent(session_id, subagent_thread, acp_thread)
1805    }
1806
1807    pub(crate) fn resume_subagent_thread(
1808        &self,
1809        session_id: acp::SessionId,
1810        cx: &mut App,
1811    ) -> Result<Rc<dyn SubagentHandle>> {
1812        let (subagent_thread, acp_thread) = self.agent.update(cx, |agent, _cx| {
1813            let session = agent
1814                .sessions
1815                .get(&session_id)
1816                .ok_or_else(|| anyhow!("No subagent session found with id {session_id}"))?;
1817            anyhow::Ok((session.thread.clone(), session.acp_thread.clone()))
1818        })??;
1819
1820        let depth = subagent_thread.read(cx).depth();
1821
1822        if let Some(parent_thread_entity) = self.thread.upgrade() {
1823            telemetry::event!(
1824                "Subagent Started",
1825                session = parent_thread_entity.read(cx).id().to_string(),
1826                subagent_session = session_id.to_string(),
1827                depth,
1828                is_resumed = true,
1829            );
1830        }
1831
1832        self.prompt_subagent(session_id, subagent_thread, acp_thread)
1833    }
1834
1835    fn prompt_subagent(
1836        &self,
1837        session_id: acp::SessionId,
1838        subagent_thread: Entity<Thread>,
1839        acp_thread: Entity<acp_thread::AcpThread>,
1840    ) -> Result<Rc<dyn SubagentHandle>> {
1841        let Some(parent_thread_entity) = self.thread.upgrade() else {
1842            anyhow::bail!("Parent thread no longer exists".to_string());
1843        };
1844        Ok(Rc::new(NativeSubagentHandle::new(
1845            session_id,
1846            subagent_thread,
1847            acp_thread,
1848            parent_thread_entity,
1849        )) as _)
1850    }
1851}
1852
1853impl ThreadEnvironment for NativeThreadEnvironment {
1854    fn create_terminal(
1855        &self,
1856        command: String,
1857        cwd: Option<PathBuf>,
1858        output_byte_limit: Option<u64>,
1859        cx: &mut AsyncApp,
1860    ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1861        let task = self.acp_thread.update(cx, |thread, cx| {
1862            thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1863        });
1864
1865        let acp_thread = self.acp_thread.clone();
1866        cx.spawn(async move |cx| {
1867            let terminal = task?.await?;
1868
1869            let (drop_tx, drop_rx) = oneshot::channel();
1870            let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone());
1871
1872            cx.spawn(async move |cx| {
1873                drop_rx.await.ok();
1874                acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1875            })
1876            .detach();
1877
1878            let handle = AcpTerminalHandle {
1879                terminal,
1880                _drop_tx: Some(drop_tx),
1881            };
1882
1883            Ok(Rc::new(handle) as _)
1884        })
1885    }
1886
1887    fn create_subagent(&self, label: String, cx: &mut App) -> Result<Rc<dyn SubagentHandle>> {
1888        self.create_subagent_thread(label, cx)
1889    }
1890
1891    fn resume_subagent(
1892        &self,
1893        session_id: acp::SessionId,
1894        cx: &mut App,
1895    ) -> Result<Rc<dyn SubagentHandle>> {
1896        self.resume_subagent_thread(session_id, cx)
1897    }
1898}
1899
1900#[derive(Debug, Clone)]
1901enum SubagentPromptResult {
1902    Completed,
1903    Cancelled,
1904    ContextWindowWarning,
1905    Error(String),
1906}
1907
1908pub struct NativeSubagentHandle {
1909    session_id: acp::SessionId,
1910    parent_thread: WeakEntity<Thread>,
1911    subagent_thread: Entity<Thread>,
1912    acp_thread: Entity<acp_thread::AcpThread>,
1913}
1914
1915impl NativeSubagentHandle {
1916    fn new(
1917        session_id: acp::SessionId,
1918        subagent_thread: Entity<Thread>,
1919        acp_thread: Entity<acp_thread::AcpThread>,
1920        parent_thread_entity: Entity<Thread>,
1921    ) -> Self {
1922        NativeSubagentHandle {
1923            session_id,
1924            subagent_thread,
1925            parent_thread: parent_thread_entity.downgrade(),
1926            acp_thread,
1927        }
1928    }
1929}
1930
1931impl SubagentHandle for NativeSubagentHandle {
1932    fn id(&self) -> acp::SessionId {
1933        self.session_id.clone()
1934    }
1935
1936    fn num_entries(&self, cx: &App) -> usize {
1937        self.acp_thread.read(cx).entries().len()
1938    }
1939
1940    fn send(&self, message: String, cx: &AsyncApp) -> Task<Result<String>> {
1941        let thread = self.subagent_thread.clone();
1942        let acp_thread = self.acp_thread.clone();
1943        let subagent_session_id = self.session_id.clone();
1944        let parent_thread = self.parent_thread.clone();
1945
1946        cx.spawn(async move |cx| {
1947            let (task, _subscription) = cx.update(|cx| {
1948                let ratio_before_prompt = thread
1949                    .read(cx)
1950                    .latest_token_usage()
1951                    .map(|usage| usage.ratio());
1952
1953                parent_thread
1954                    .update(cx, |parent_thread, _cx| {
1955                        parent_thread.register_running_subagent(thread.downgrade())
1956                    })
1957                    .ok();
1958
1959                let task = acp_thread.update(cx, |acp_thread, cx| {
1960                    acp_thread.send(vec![message.into()], cx)
1961                });
1962
1963                let (token_limit_tx, token_limit_rx) = oneshot::channel::<()>();
1964                let mut token_limit_tx = Some(token_limit_tx);
1965
1966                let subscription = cx.subscribe(
1967                    &thread,
1968                    move |_thread, event: &TokenUsageUpdated, _cx| {
1969                        if let Some(usage) = &event.0 {
1970                            let old_ratio = ratio_before_prompt
1971                                .clone()
1972                                .unwrap_or(TokenUsageRatio::Normal);
1973                            let new_ratio = usage.ratio();
1974                            if old_ratio == TokenUsageRatio::Normal
1975                                && new_ratio == TokenUsageRatio::Warning
1976                            {
1977                                if let Some(tx) = token_limit_tx.take() {
1978                                    tx.send(()).ok();
1979                                }
1980                            }
1981                        }
1982                    },
1983                );
1984
1985                let wait_for_prompt = cx
1986                    .background_spawn(async move {
1987                        futures::select! {
1988                            response = task.fuse() => match response {
1989                                Ok(Some(response)) => {
1990                                    match response.stop_reason {
1991                                        acp::StopReason::Cancelled => SubagentPromptResult::Cancelled,
1992                                        acp::StopReason::MaxTokens => SubagentPromptResult::Error("The agent reached the maximum number of tokens.".into()),
1993                                        acp::StopReason::MaxTurnRequests => SubagentPromptResult::Error("The agent reached the maximum number of allowed requests between user turns. Try prompting again.".into()),
1994                                        acp::StopReason::Refusal => SubagentPromptResult::Error("The agent refused to process that prompt. Try again.".into()),
1995                                        acp::StopReason::EndTurn | _ => SubagentPromptResult::Completed,
1996                                    }
1997                                }
1998                                Ok(None) => SubagentPromptResult::Error("No response from the agent. You can try messaging again.".into()),
1999                                Err(error) => SubagentPromptResult::Error(error.to_string()),
2000                            },
2001                            _ = token_limit_rx.fuse() => SubagentPromptResult::ContextWindowWarning,
2002                        }
2003                    });
2004
2005                (wait_for_prompt, subscription)
2006            });
2007
2008            let result = match task.await {
2009                SubagentPromptResult::Completed => thread.read_with(cx, |thread, _cx| {
2010                    thread
2011                        .last_message()
2012                        .and_then(|message| {
2013                            let content = message.as_agent_message()?
2014                                .content
2015                                .iter()
2016                                .filter_map(|c| match c {
2017                                    AgentMessageContent::Text(text) => Some(text.as_str()),
2018                                    _ => None,
2019                                })
2020                                .join("\n\n");
2021                            if content.is_empty() {
2022                                None
2023                            } else {
2024                                Some( content)
2025                            }
2026                        })
2027                        .context("No response from subagent")
2028                }),
2029                SubagentPromptResult::Cancelled => Err(anyhow!("User canceled")),
2030                SubagentPromptResult::Error(message) => Err(anyhow!("{message}")),
2031                SubagentPromptResult::ContextWindowWarning => {
2032                    thread.update(cx, |thread, cx| thread.cancel(cx)).await;
2033                    Err(anyhow!(
2034                        "The agent is nearing the end of its context window and has been \
2035                         stopped. You can prompt the thread again to have the agent wrap up \
2036                         or hand off its work."
2037                    ))
2038                }
2039            };
2040
2041            parent_thread
2042                .update(cx, |parent_thread, cx| {
2043                    parent_thread.unregister_running_subagent(&subagent_session_id, cx)
2044                })
2045                .ok();
2046
2047            result
2048        })
2049    }
2050}
2051
2052pub struct AcpTerminalHandle {
2053    terminal: Entity<acp_thread::Terminal>,
2054    _drop_tx: Option<oneshot::Sender<()>>,
2055}
2056
2057impl TerminalHandle for AcpTerminalHandle {
2058    fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
2059        Ok(self.terminal.read_with(cx, |term, _cx| term.id().clone()))
2060    }
2061
2062    fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
2063        Ok(self
2064            .terminal
2065            .read_with(cx, |term, _cx| term.wait_for_exit()))
2066    }
2067
2068    fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
2069        Ok(self
2070            .terminal
2071            .read_with(cx, |term, cx| term.current_output(cx)))
2072    }
2073
2074    fn kill(&self, cx: &AsyncApp) -> Result<()> {
2075        cx.update(|cx| {
2076            self.terminal.update(cx, |terminal, cx| {
2077                terminal.kill(cx);
2078            });
2079        });
2080        Ok(())
2081    }
2082
2083    fn was_stopped_by_user(&self, cx: &AsyncApp) -> Result<bool> {
2084        Ok(self
2085            .terminal
2086            .read_with(cx, |term, _cx| term.was_stopped_by_user()))
2087    }
2088}
2089
2090#[cfg(test)]
2091mod internal_tests {
2092    use std::path::Path;
2093
2094    use super::*;
2095    use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
2096    use fs::FakeFs;
2097    use gpui::TestAppContext;
2098    use indoc::formatdoc;
2099    use language_model::fake_provider::{FakeLanguageModel, FakeLanguageModelProvider};
2100    use language_model::{
2101        LanguageModelCompletionEvent, LanguageModelProviderId, LanguageModelProviderName,
2102    };
2103    use serde_json::json;
2104    use settings::SettingsStore;
2105    use util::{path, rel_path::rel_path};
2106
2107    #[gpui::test]
2108    async fn test_maintaining_project_context(cx: &mut TestAppContext) {
2109        init_test(cx);
2110        let fs = FakeFs::new(cx.executor());
2111        fs.insert_tree(
2112            "/",
2113            json!({
2114                "a": {}
2115            }),
2116        )
2117        .await;
2118        let project = Project::test(fs.clone(), [], cx).await;
2119        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2120        let agent =
2121            cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2122
2123        // Creating a session registers the project and triggers context building.
2124        let connection = NativeAgentConnection(agent.clone());
2125        let _acp_thread = cx
2126            .update(|cx| {
2127                Rc::new(connection).new_session(
2128                    project.clone(),
2129                    PathList::new(&[Path::new("/")]),
2130                    cx,
2131                )
2132            })
2133            .await
2134            .unwrap();
2135        cx.run_until_parked();
2136
2137        let thread = agent.read_with(cx, |agent, _cx| {
2138            agent.sessions.values().next().unwrap().thread.clone()
2139        });
2140
2141        agent.read_with(cx, |agent, cx| {
2142            let project_id = project.entity_id();
2143            let state = agent.projects.get(&project_id).unwrap();
2144            assert_eq!(state.project_context.read(cx).worktrees, vec![]);
2145            assert_eq!(thread.read(cx).project_context().read(cx).worktrees, vec![]);
2146        });
2147
2148        let worktree = project
2149            .update(cx, |project, cx| project.create_worktree("/a", true, cx))
2150            .await
2151            .unwrap();
2152        cx.run_until_parked();
2153        agent.read_with(cx, |agent, cx| {
2154            let project_id = project.entity_id();
2155            let state = agent.projects.get(&project_id).unwrap();
2156            let expected_worktrees = vec![WorktreeContext {
2157                root_name: "a".into(),
2158                abs_path: Path::new("/a").into(),
2159                rules_file: None,
2160            }];
2161            assert_eq!(state.project_context.read(cx).worktrees, expected_worktrees);
2162            assert_eq!(
2163                thread.read(cx).project_context().read(cx).worktrees,
2164                expected_worktrees
2165            );
2166        });
2167
2168        // Creating `/a/.rules` updates the project context.
2169        fs.insert_file("/a/.rules", Vec::new()).await;
2170        cx.run_until_parked();
2171        agent.read_with(cx, |agent, cx| {
2172            let project_id = project.entity_id();
2173            let state = agent.projects.get(&project_id).unwrap();
2174            let rules_entry = worktree
2175                .read(cx)
2176                .entry_for_path(rel_path(".rules"))
2177                .unwrap();
2178            let expected_worktrees = vec![WorktreeContext {
2179                root_name: "a".into(),
2180                abs_path: Path::new("/a").into(),
2181                rules_file: Some(RulesFileContext {
2182                    path_in_worktree: rel_path(".rules").into(),
2183                    text: "".into(),
2184                    project_entry_id: rules_entry.id.to_usize(),
2185                }),
2186            }];
2187            assert_eq!(state.project_context.read(cx).worktrees, expected_worktrees);
2188            assert_eq!(
2189                thread.read(cx).project_context().read(cx).worktrees,
2190                expected_worktrees
2191            );
2192        });
2193    }
2194
2195    #[gpui::test]
2196    async fn test_listing_models(cx: &mut TestAppContext) {
2197        init_test(cx);
2198        let fs = FakeFs::new(cx.executor());
2199        fs.insert_tree("/", json!({ "a": {}  })).await;
2200        let project = Project::test(fs.clone(), [], cx).await;
2201        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2202        let connection =
2203            NativeAgentConnection(cx.update(|cx| {
2204                NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx)
2205            }));
2206
2207        // Create a thread/session
2208        let acp_thread = cx
2209            .update(|cx| {
2210                Rc::new(connection.clone()).new_session(
2211                    project.clone(),
2212                    PathList::new(&[Path::new("/a")]),
2213                    cx,
2214                )
2215            })
2216            .await
2217            .unwrap();
2218
2219        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2220
2221        let models = cx
2222            .update(|cx| {
2223                connection
2224                    .model_selector(&session_id)
2225                    .unwrap()
2226                    .list_models(cx)
2227            })
2228            .await
2229            .unwrap();
2230
2231        let acp_thread::AgentModelList::Grouped(models) = models else {
2232            panic!("Unexpected model group");
2233        };
2234        assert_eq!(
2235            models,
2236            IndexMap::from_iter([(
2237                AgentModelGroupName("Fake".into()),
2238                vec![AgentModelInfo {
2239                    id: acp::ModelId::new("fake/fake"),
2240                    name: "Fake".into(),
2241                    description: None,
2242                    icon: Some(acp_thread::AgentModelIcon::Named(
2243                        ui::IconName::ZedAssistant
2244                    )),
2245                    is_latest: false,
2246                    cost: None,
2247                }]
2248            )])
2249        );
2250    }
2251
2252    #[gpui::test]
2253    async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
2254        init_test(cx);
2255        let fs = FakeFs::new(cx.executor());
2256        fs.create_dir(paths::settings_file().parent().unwrap())
2257            .await
2258            .unwrap();
2259        fs.insert_file(
2260            paths::settings_file(),
2261            json!({
2262                "agent": {
2263                    "default_model": {
2264                        "provider": "foo",
2265                        "model": "bar"
2266                    }
2267                }
2268            })
2269            .to_string()
2270            .into_bytes(),
2271        )
2272        .await;
2273        let project = Project::test(fs.clone(), [], cx).await;
2274
2275        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2276
2277        // Create the agent and connection
2278        let agent =
2279            cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2280        let connection = NativeAgentConnection(agent.clone());
2281
2282        // Create a thread/session
2283        let acp_thread = cx
2284            .update(|cx| {
2285                Rc::new(connection.clone()).new_session(
2286                    project.clone(),
2287                    PathList::new(&[Path::new("/a")]),
2288                    cx,
2289                )
2290            })
2291            .await
2292            .unwrap();
2293
2294        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2295
2296        // Select a model
2297        let selector = connection.model_selector(&session_id).unwrap();
2298        let model_id = acp::ModelId::new("fake/fake");
2299        cx.update(|cx| selector.select_model(model_id.clone(), cx))
2300            .await
2301            .unwrap();
2302
2303        // Verify the thread has the selected model
2304        agent.read_with(cx, |agent, _| {
2305            let session = agent.sessions.get(&session_id).unwrap();
2306            session.thread.read_with(cx, |thread, _| {
2307                assert_eq!(thread.model().unwrap().id().0, "fake");
2308            });
2309        });
2310
2311        cx.run_until_parked();
2312
2313        // Verify settings file was updated
2314        let settings_content = fs.load(paths::settings_file()).await.unwrap();
2315        let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2316
2317        // Check that the agent settings contain the selected model
2318        assert_eq!(
2319            settings_json["agent"]["default_model"]["model"],
2320            json!("fake")
2321        );
2322        assert_eq!(
2323            settings_json["agent"]["default_model"]["provider"],
2324            json!("fake")
2325        );
2326
2327        // Register a thinking model and select it.
2328        cx.update(|cx| {
2329            let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2330                "fake-corp",
2331                "fake-thinking",
2332                "Fake Thinking",
2333                true,
2334            ));
2335            let thinking_provider = Arc::new(
2336                FakeLanguageModelProvider::new(
2337                    LanguageModelProviderId::from("fake-corp".to_string()),
2338                    LanguageModelProviderName::from("Fake Corp".to_string()),
2339                )
2340                .with_models(vec![thinking_model]),
2341            );
2342            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2343                registry.register_provider(thinking_provider, cx);
2344            });
2345        });
2346        agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2347
2348        let selector = connection.model_selector(&session_id).unwrap();
2349        cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2350            .await
2351            .unwrap();
2352        cx.run_until_parked();
2353
2354        // Verify enable_thinking was written to settings as true.
2355        let settings_content = fs.load(paths::settings_file()).await.unwrap();
2356        let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
2357        assert_eq!(
2358            settings_json["agent"]["default_model"]["enable_thinking"],
2359            json!(true),
2360            "selecting a thinking model should persist enable_thinking: true to settings"
2361        );
2362    }
2363
2364    #[gpui::test]
2365    async fn test_select_model_updates_thinking_enabled(cx: &mut TestAppContext) {
2366        init_test(cx);
2367        let fs = FakeFs::new(cx.executor());
2368        fs.create_dir(paths::settings_file().parent().unwrap())
2369            .await
2370            .unwrap();
2371        fs.insert_file(paths::settings_file(), b"{}".to_vec()).await;
2372        let project = Project::test(fs.clone(), [], cx).await;
2373
2374        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2375        let agent =
2376            cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), None, fs.clone(), cx));
2377        let connection = NativeAgentConnection(agent.clone());
2378
2379        let acp_thread = cx
2380            .update(|cx| {
2381                Rc::new(connection.clone()).new_session(
2382                    project.clone(),
2383                    PathList::new(&[Path::new("/a")]),
2384                    cx,
2385                )
2386            })
2387            .await
2388            .unwrap();
2389        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
2390
2391        // Register a second provider with a thinking model.
2392        cx.update(|cx| {
2393            let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2394                "fake-corp",
2395                "fake-thinking",
2396                "Fake Thinking",
2397                true,
2398            ));
2399            let thinking_provider = Arc::new(
2400                FakeLanguageModelProvider::new(
2401                    LanguageModelProviderId::from("fake-corp".to_string()),
2402                    LanguageModelProviderName::from("Fake Corp".to_string()),
2403                )
2404                .with_models(vec![thinking_model]),
2405            );
2406            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2407                registry.register_provider(thinking_provider, cx);
2408            });
2409        });
2410        // Refresh the agent's model list so it picks up the new provider.
2411        agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2412
2413        // Thread starts with thinking_enabled = false (the default).
2414        agent.read_with(cx, |agent, _| {
2415            let session = agent.sessions.get(&session_id).unwrap();
2416            session.thread.read_with(cx, |thread, _| {
2417                assert!(!thread.thinking_enabled(), "thinking defaults to false");
2418            });
2419        });
2420
2421        // Select the thinking model via select_model.
2422        let selector = connection.model_selector(&session_id).unwrap();
2423        cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2424            .await
2425            .unwrap();
2426
2427        // select_model should have enabled thinking based on the model's supports_thinking().
2428        agent.read_with(cx, |agent, _| {
2429            let session = agent.sessions.get(&session_id).unwrap();
2430            session.thread.read_with(cx, |thread, _| {
2431                assert!(
2432                    thread.thinking_enabled(),
2433                    "select_model should enable thinking when model supports it"
2434                );
2435            });
2436        });
2437
2438        // Switch back to the non-thinking model.
2439        let selector = connection.model_selector(&session_id).unwrap();
2440        cx.update(|cx| selector.select_model(acp::ModelId::new("fake/fake"), cx))
2441            .await
2442            .unwrap();
2443
2444        // select_model should have disabled thinking.
2445        agent.read_with(cx, |agent, _| {
2446            let session = agent.sessions.get(&session_id).unwrap();
2447            session.thread.read_with(cx, |thread, _| {
2448                assert!(
2449                    !thread.thinking_enabled(),
2450                    "select_model should disable thinking when model does not support it"
2451                );
2452            });
2453        });
2454    }
2455
2456    #[gpui::test]
2457    async fn test_loaded_thread_preserves_thinking_enabled(cx: &mut TestAppContext) {
2458        init_test(cx);
2459        let fs = FakeFs::new(cx.executor());
2460        fs.insert_tree("/", json!({ "a": {} })).await;
2461        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2462        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2463        let agent = cx.update(|cx| {
2464            NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2465        });
2466        let connection = Rc::new(NativeAgentConnection(agent.clone()));
2467
2468        // Register a thinking model.
2469        let thinking_model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2470            "fake-corp",
2471            "fake-thinking",
2472            "Fake Thinking",
2473            true,
2474        ));
2475        let thinking_provider = Arc::new(
2476            FakeLanguageModelProvider::new(
2477                LanguageModelProviderId::from("fake-corp".to_string()),
2478                LanguageModelProviderName::from("Fake Corp".to_string()),
2479            )
2480            .with_models(vec![thinking_model.clone()]),
2481        );
2482        cx.update(|cx| {
2483            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2484                registry.register_provider(thinking_provider, cx);
2485            });
2486        });
2487        agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2488
2489        // Create a thread and select the thinking model.
2490        let acp_thread = cx
2491            .update(|cx| {
2492                connection.clone().new_session(
2493                    project.clone(),
2494                    PathList::new(&[Path::new("/a")]),
2495                    cx,
2496                )
2497            })
2498            .await
2499            .unwrap();
2500        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2501
2502        let selector = connection.model_selector(&session_id).unwrap();
2503        cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/fake-thinking"), cx))
2504            .await
2505            .unwrap();
2506
2507        // Verify thinking is enabled after selecting the thinking model.
2508        let thread = agent.read_with(cx, |agent, _| {
2509            agent.sessions.get(&session_id).unwrap().thread.clone()
2510        });
2511        thread.read_with(cx, |thread, _| {
2512            assert!(
2513                thread.thinking_enabled(),
2514                "thinking should be enabled after selecting thinking model"
2515            );
2516        });
2517
2518        // Send a message so the thread gets persisted.
2519        let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2520        let send = cx.foreground_executor().spawn(send);
2521        cx.run_until_parked();
2522
2523        thinking_model.send_last_completion_stream_text_chunk("Response.");
2524        thinking_model.end_last_completion_stream();
2525
2526        send.await.unwrap();
2527        cx.run_until_parked();
2528
2529        // Close the session so it can be reloaded from disk.
2530        cx.update(|cx| connection.clone().close_session(&session_id, cx))
2531            .await
2532            .unwrap();
2533        drop(thread);
2534        drop(acp_thread);
2535        agent.read_with(cx, |agent, _| {
2536            assert!(agent.sessions.is_empty());
2537        });
2538
2539        // Reload the thread and verify thinking_enabled is still true.
2540        let reloaded_acp_thread = agent
2541            .update(cx, |agent, cx| {
2542                agent.open_thread(session_id.clone(), project.clone(), cx)
2543            })
2544            .await
2545            .unwrap();
2546        let reloaded_thread = agent.read_with(cx, |agent, _| {
2547            agent.sessions.get(&session_id).unwrap().thread.clone()
2548        });
2549        reloaded_thread.read_with(cx, |thread, _| {
2550            assert!(
2551                thread.thinking_enabled(),
2552                "thinking_enabled should be preserved when reloading a thread with a thinking model"
2553            );
2554        });
2555
2556        drop(reloaded_acp_thread);
2557    }
2558
2559    #[gpui::test]
2560    async fn test_loaded_thread_preserves_model(cx: &mut TestAppContext) {
2561        init_test(cx);
2562        let fs = FakeFs::new(cx.executor());
2563        fs.insert_tree("/", json!({ "a": {} })).await;
2564        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2565        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2566        let agent = cx.update(|cx| {
2567            NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2568        });
2569        let connection = Rc::new(NativeAgentConnection(agent.clone()));
2570
2571        // Register a model where id() != name(), like real Anthropic models
2572        // (e.g. id="claude-sonnet-4-5-thinking-latest", name="Claude Sonnet 4.5 Thinking").
2573        let model = Arc::new(FakeLanguageModel::with_id_and_thinking(
2574            "fake-corp",
2575            "custom-model-id",
2576            "Custom Model Display Name",
2577            false,
2578        ));
2579        let provider = Arc::new(
2580            FakeLanguageModelProvider::new(
2581                LanguageModelProviderId::from("fake-corp".to_string()),
2582                LanguageModelProviderName::from("Fake Corp".to_string()),
2583            )
2584            .with_models(vec![model.clone()]),
2585        );
2586        cx.update(|cx| {
2587            LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
2588                registry.register_provider(provider, cx);
2589            });
2590        });
2591        agent.update(cx, |agent, cx| agent.models.refresh_list(cx));
2592
2593        // Create a thread and select the model.
2594        let acp_thread = cx
2595            .update(|cx| {
2596                connection.clone().new_session(
2597                    project.clone(),
2598                    PathList::new(&[Path::new("/a")]),
2599                    cx,
2600                )
2601            })
2602            .await
2603            .unwrap();
2604        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2605
2606        let selector = connection.model_selector(&session_id).unwrap();
2607        cx.update(|cx| selector.select_model(acp::ModelId::new("fake-corp/custom-model-id"), cx))
2608            .await
2609            .unwrap();
2610
2611        let thread = agent.read_with(cx, |agent, _| {
2612            agent.sessions.get(&session_id).unwrap().thread.clone()
2613        });
2614        thread.read_with(cx, |thread, _| {
2615            assert_eq!(
2616                thread.model().unwrap().id().0.as_ref(),
2617                "custom-model-id",
2618                "model should be set before persisting"
2619            );
2620        });
2621
2622        // Send a message so the thread gets persisted.
2623        let send = acp_thread.update(cx, |thread, cx| thread.send(vec!["Hello".into()], cx));
2624        let send = cx.foreground_executor().spawn(send);
2625        cx.run_until_parked();
2626
2627        model.send_last_completion_stream_text_chunk("Response.");
2628        model.end_last_completion_stream();
2629
2630        send.await.unwrap();
2631        cx.run_until_parked();
2632
2633        // Close the session so it can be reloaded from disk.
2634        cx.update(|cx| connection.clone().close_session(&session_id, cx))
2635            .await
2636            .unwrap();
2637        drop(thread);
2638        drop(acp_thread);
2639        agent.read_with(cx, |agent, _| {
2640            assert!(agent.sessions.is_empty());
2641        });
2642
2643        // Reload the thread and verify the model was preserved.
2644        let reloaded_acp_thread = agent
2645            .update(cx, |agent, cx| {
2646                agent.open_thread(session_id.clone(), project.clone(), cx)
2647            })
2648            .await
2649            .unwrap();
2650        let reloaded_thread = agent.read_with(cx, |agent, _| {
2651            agent.sessions.get(&session_id).unwrap().thread.clone()
2652        });
2653        reloaded_thread.read_with(cx, |thread, _| {
2654            let reloaded_model = thread
2655                .model()
2656                .expect("model should be present after reload");
2657            assert_eq!(
2658                reloaded_model.id().0.as_ref(),
2659                "custom-model-id",
2660                "reloaded thread should have the same model, not fall back to the default"
2661            );
2662        });
2663
2664        drop(reloaded_acp_thread);
2665    }
2666
2667    #[gpui::test]
2668    async fn test_save_load_thread(cx: &mut TestAppContext) {
2669        init_test(cx);
2670        let fs = FakeFs::new(cx.executor());
2671        fs.insert_tree(
2672            "/",
2673            json!({
2674                "a": {
2675                    "b.md": "Lorem"
2676                }
2677            }),
2678        )
2679        .await;
2680        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
2681        let thread_store = cx.new(|cx| ThreadStore::new(cx));
2682        let agent = cx.update(|cx| {
2683            NativeAgent::new(thread_store.clone(), Templates::new(), None, fs.clone(), cx)
2684        });
2685        let connection = Rc::new(NativeAgentConnection(agent.clone()));
2686
2687        let acp_thread = cx
2688            .update(|cx| {
2689                connection
2690                    .clone()
2691                    .new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
2692            })
2693            .await
2694            .unwrap();
2695        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
2696        let thread = agent.read_with(cx, |agent, _| {
2697            agent.sessions.get(&session_id).unwrap().thread.clone()
2698        });
2699
2700        // Ensure empty threads are not saved, even if they get mutated.
2701        let model = Arc::new(FakeLanguageModel::default());
2702        let summary_model = Arc::new(FakeLanguageModel::default());
2703        thread.update(cx, |thread, cx| {
2704            thread.set_model(model.clone(), cx);
2705            thread.set_summarization_model(Some(summary_model.clone()), cx);
2706        });
2707        cx.run_until_parked();
2708        assert_eq!(thread_entries(&thread_store, cx), vec![]);
2709
2710        let send = acp_thread.update(cx, |thread, cx| {
2711            thread.send(
2712                vec![
2713                    "What does ".into(),
2714                    acp::ContentBlock::ResourceLink(acp::ResourceLink::new(
2715                        "b.md",
2716                        MentionUri::File {
2717                            abs_path: path!("/a/b.md").into(),
2718                        }
2719                        .to_uri()
2720                        .to_string(),
2721                    )),
2722                    " mean?".into(),
2723                ],
2724                cx,
2725            )
2726        });
2727        let send = cx.foreground_executor().spawn(send);
2728        cx.run_until_parked();
2729
2730        model.send_last_completion_stream_text_chunk("Lorem.");
2731        model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
2732            language_model::TokenUsage {
2733                input_tokens: 150,
2734                output_tokens: 75,
2735                ..Default::default()
2736            },
2737        ));
2738        model.end_last_completion_stream();
2739        cx.run_until_parked();
2740        summary_model
2741            .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
2742        summary_model.end_last_completion_stream();
2743
2744        send.await.unwrap();
2745        let uri = MentionUri::File {
2746            abs_path: path!("/a/b.md").into(),
2747        }
2748        .to_uri();
2749        acp_thread.read_with(cx, |thread, cx| {
2750            assert_eq!(
2751                thread.to_markdown(cx),
2752                formatdoc! {"
2753                    ## User
2754
2755                    What does [@b.md]({uri}) mean?
2756
2757                    ## Assistant
2758
2759                    Lorem.
2760
2761                "}
2762            )
2763        });
2764
2765        cx.run_until_parked();
2766
2767        // Set a draft prompt with rich content blocks before saving.
2768        let draft_blocks = vec![
2769            acp::ContentBlock::Text(acp::TextContent::new("Check out ")),
2770            acp::ContentBlock::ResourceLink(acp::ResourceLink::new("b.md", uri.to_string())),
2771            acp::ContentBlock::Text(acp::TextContent::new(" please")),
2772        ];
2773        acp_thread.update(cx, |thread, _cx| {
2774            thread.set_draft_prompt(Some(draft_blocks.clone()));
2775        });
2776        thread.update(cx, |thread, _cx| {
2777            thread.set_ui_scroll_position(Some(gpui::ListOffset {
2778                item_ix: 5,
2779                offset_in_item: gpui::px(12.5),
2780            }));
2781        });
2782        thread.update(cx, |_thread, cx| cx.notify());
2783        cx.run_until_parked();
2784
2785        // Close the session so it can be reloaded from disk.
2786        cx.update(|cx| connection.clone().close_session(&session_id, cx))
2787            .await
2788            .unwrap();
2789        drop(thread);
2790        drop(acp_thread);
2791        agent.read_with(cx, |agent, _| {
2792            assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
2793        });
2794
2795        // Ensure the thread can be reloaded from disk.
2796        assert_eq!(
2797            thread_entries(&thread_store, cx),
2798            vec![(
2799                session_id.clone(),
2800                format!("Explaining {}", path!("/a/b.md"))
2801            )]
2802        );
2803        let acp_thread = agent
2804            .update(cx, |agent, cx| {
2805                agent.open_thread(session_id.clone(), project.clone(), cx)
2806            })
2807            .await
2808            .unwrap();
2809        acp_thread.read_with(cx, |thread, cx| {
2810            assert_eq!(
2811                thread.to_markdown(cx),
2812                formatdoc! {"
2813                    ## User
2814
2815                    What does [@b.md]({uri}) mean?
2816
2817                    ## Assistant
2818
2819                    Lorem.
2820
2821                "}
2822            )
2823        });
2824
2825        // Ensure the draft prompt with rich content blocks survived the round-trip.
2826        acp_thread.read_with(cx, |thread, _| {
2827            assert_eq!(thread.draft_prompt(), Some(draft_blocks.as_slice()));
2828        });
2829
2830        // Ensure token usage survived the round-trip.
2831        acp_thread.read_with(cx, |thread, _| {
2832            let usage = thread
2833                .token_usage()
2834                .expect("token usage should be restored after reload");
2835            assert_eq!(usage.input_tokens, 150);
2836            assert_eq!(usage.output_tokens, 75);
2837        });
2838
2839        // Ensure scroll position survived the round-trip.
2840        acp_thread.read_with(cx, |thread, _| {
2841            let scroll = thread
2842                .ui_scroll_position()
2843                .expect("scroll position should be restored after reload");
2844            assert_eq!(scroll.item_ix, 5);
2845            assert_eq!(scroll.offset_in_item, gpui::px(12.5));
2846        });
2847    }
2848
2849    fn thread_entries(
2850        thread_store: &Entity<ThreadStore>,
2851        cx: &mut TestAppContext,
2852    ) -> Vec<(acp::SessionId, String)> {
2853        thread_store.read_with(cx, |store, _| {
2854            store
2855                .entries()
2856                .map(|entry| (entry.id.clone(), entry.title.to_string()))
2857                .collect::<Vec<_>>()
2858        })
2859    }
2860
2861    fn init_test(cx: &mut TestAppContext) {
2862        env_logger::try_init().ok();
2863        cx.update(|cx| {
2864            let settings_store = SettingsStore::test(cx);
2865            cx.set_global(settings_store);
2866
2867            LanguageModelRegistry::test(cx);
2868        });
2869    }
2870}
2871
2872fn mcp_message_content_to_acp_content_block(
2873    content: context_server::types::MessageContent,
2874) -> acp::ContentBlock {
2875    match content {
2876        context_server::types::MessageContent::Text {
2877            text,
2878            annotations: _,
2879        } => text.into(),
2880        context_server::types::MessageContent::Image {
2881            data,
2882            mime_type,
2883            annotations: _,
2884        } => acp::ContentBlock::Image(acp::ImageContent::new(data, mime_type)),
2885        context_server::types::MessageContent::Audio {
2886            data,
2887            mime_type,
2888            annotations: _,
2889        } => acp::ContentBlock::Audio(acp::AudioContent::new(data, mime_type)),
2890        context_server::types::MessageContent::Resource {
2891            resource,
2892            annotations: _,
2893        } => {
2894            let mut link =
2895                acp::ResourceLink::new(resource.uri.to_string(), resource.uri.to_string());
2896            if let Some(mime_type) = resource.mime_type {
2897                link = link.mime_type(mime_type);
2898            }
2899            acp::ContentBlock::ResourceLink(link)
2900        }
2901    }
2902}