agent.rs

   1mod db;
   2mod edit_agent;
   3mod history_store;
   4mod legacy_thread;
   5mod native_agent_server;
   6pub mod outline;
   7mod templates;
   8mod thread;
   9mod tools;
  10
  11#[cfg(test)]
  12mod tests;
  13
  14pub use db::*;
  15pub use history_store::*;
  16pub use native_agent_server::NativeAgentServer;
  17pub use templates::*;
  18pub use thread::*;
  19pub use tools::*;
  20
  21use acp_thread::{AcpThread, AgentModelSelector};
  22use agent_client_protocol as acp;
  23use anyhow::{Context as _, Result, anyhow};
  24use chrono::{DateTime, Utc};
  25use collections::{HashSet, IndexMap};
  26use fs::Fs;
  27use futures::channel::{mpsc, oneshot};
  28use futures::future::Shared;
  29use futures::{StreamExt, future};
  30use gpui::{
  31    App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,
  32};
  33use language_model::{LanguageModel, LanguageModelProvider, LanguageModelRegistry};
  34use project::{Project, ProjectItem, ProjectPath, Worktree};
  35use prompt_store::{
  36    ProjectContext, PromptStore, RulesFileContext, UserRulesContext, WorktreeContext,
  37};
  38use serde::{Deserialize, Serialize};
  39use settings::{LanguageModelSelection, update_settings_file};
  40use std::any::Any;
  41use std::collections::HashMap;
  42use std::path::{Path, PathBuf};
  43use std::rc::Rc;
  44use std::sync::Arc;
  45use util::ResultExt;
  46use util::rel_path::RelPath;
  47
  48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
  49pub struct ProjectSnapshot {
  50    pub worktree_snapshots: Vec<project::telemetry_snapshot::TelemetryWorktreeSnapshot>,
  51    pub timestamp: DateTime<Utc>,
  52}
  53
  54const RULES_FILE_NAMES: [&str; 9] = [
  55    ".rules",
  56    ".cursorrules",
  57    ".windsurfrules",
  58    ".clinerules",
  59    ".github/copilot-instructions.md",
  60    "CLAUDE.md",
  61    "AGENT.md",
  62    "AGENTS.md",
  63    "GEMINI.md",
  64];
  65
  66pub struct RulesLoadingError {
  67    pub message: SharedString,
  68}
  69
  70/// Holds both the internal Thread and the AcpThread for a session
  71struct Session {
  72    /// The internal thread that processes messages
  73    thread: Entity<Thread>,
  74    /// The ACP thread that handles protocol communication
  75    acp_thread: WeakEntity<acp_thread::AcpThread>,
  76    pending_save: Task<()>,
  77    _subscriptions: Vec<Subscription>,
  78}
  79
  80pub struct LanguageModels {
  81    /// Access language model by ID
  82    models: HashMap<acp::ModelId, Arc<dyn LanguageModel>>,
  83    /// Cached list for returning language model information
  84    model_list: acp_thread::AgentModelList,
  85    refresh_models_rx: watch::Receiver<()>,
  86    refresh_models_tx: watch::Sender<()>,
  87    _authenticate_all_providers_task: Task<()>,
  88}
  89
  90impl LanguageModels {
  91    fn new(cx: &mut App) -> Self {
  92        let (refresh_models_tx, refresh_models_rx) = watch::channel(());
  93
  94        let mut this = Self {
  95            models: HashMap::default(),
  96            model_list: acp_thread::AgentModelList::Grouped(IndexMap::default()),
  97            refresh_models_rx,
  98            refresh_models_tx,
  99            _authenticate_all_providers_task: Self::authenticate_all_language_model_providers(cx),
 100        };
 101        this.refresh_list(cx);
 102        this
 103    }
 104
 105    fn refresh_list(&mut self, cx: &App) {
 106        let providers = LanguageModelRegistry::global(cx)
 107            .read(cx)
 108            .providers()
 109            .into_iter()
 110            .filter(|provider| provider.is_authenticated(cx))
 111            .collect::<Vec<_>>();
 112
 113        let mut language_model_list = IndexMap::default();
 114        let mut recommended_models = HashSet::default();
 115
 116        let mut recommended = Vec::new();
 117        for provider in &providers {
 118            for model in provider.recommended_models(cx) {
 119                recommended_models.insert((model.provider_id(), model.id()));
 120                recommended.push(Self::map_language_model_to_info(&model, provider));
 121            }
 122        }
 123        if !recommended.is_empty() {
 124            language_model_list.insert(
 125                acp_thread::AgentModelGroupName("Recommended".into()),
 126                recommended,
 127            );
 128        }
 129
 130        let mut models = HashMap::default();
 131        for provider in providers {
 132            let mut provider_models = Vec::new();
 133            for model in provider.provided_models(cx) {
 134                let model_info = Self::map_language_model_to_info(&model, &provider);
 135                let model_id = model_info.id.clone();
 136                if !recommended_models.contains(&(model.provider_id(), model.id())) {
 137                    provider_models.push(model_info);
 138                }
 139                models.insert(model_id, model);
 140            }
 141            if !provider_models.is_empty() {
 142                language_model_list.insert(
 143                    acp_thread::AgentModelGroupName(provider.name().0.clone()),
 144                    provider_models,
 145                );
 146            }
 147        }
 148
 149        self.models = models;
 150        self.model_list = acp_thread::AgentModelList::Grouped(language_model_list);
 151        self.refresh_models_tx.send(()).ok();
 152    }
 153
 154    fn watch(&self) -> watch::Receiver<()> {
 155        self.refresh_models_rx.clone()
 156    }
 157
 158    pub fn model_from_id(&self, model_id: &acp::ModelId) -> Option<Arc<dyn LanguageModel>> {
 159        self.models.get(model_id).cloned()
 160    }
 161
 162    fn map_language_model_to_info(
 163        model: &Arc<dyn LanguageModel>,
 164        provider: &Arc<dyn LanguageModelProvider>,
 165    ) -> acp_thread::AgentModelInfo {
 166        acp_thread::AgentModelInfo {
 167            id: Self::model_id(model),
 168            name: model.name().0,
 169            description: None,
 170            icon: Some(provider.icon()),
 171        }
 172    }
 173
 174    fn model_id(model: &Arc<dyn LanguageModel>) -> acp::ModelId {
 175        acp::ModelId(format!("{}/{}", model.provider_id().0, model.id().0).into())
 176    }
 177
 178    fn authenticate_all_language_model_providers(cx: &mut App) -> Task<()> {
 179        let authenticate_all_providers = LanguageModelRegistry::global(cx)
 180            .read(cx)
 181            .providers()
 182            .iter()
 183            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))
 184            .collect::<Vec<_>>();
 185
 186        cx.background_spawn(async move {
 187            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {
 188                if let Err(err) = authenticate_task.await {
 189                    match err {
 190                        language_model::AuthenticateError::CredentialsNotFound => {
 191                            // Since we're authenticating these providers in the
 192                            // background for the purposes of populating the
 193                            // language selector, we don't care about providers
 194                            // where the credentials are not found.
 195                        }
 196                        language_model::AuthenticateError::ConnectionRefused => {
 197                            // Not logging connection refused errors as they are mostly from LM Studio's noisy auth failures.
 198                            // LM Studio only has one auth method (endpoint call) which fails for users who haven't enabled it.
 199                            // TODO: Better manage LM Studio auth logic to avoid these noisy failures.
 200                        }
 201                        _ => {
 202                            // Some providers have noisy failure states that we
 203                            // don't want to spam the logs with every time the
 204                            // language model selector is initialized.
 205                            //
 206                            // Ideally these should have more clear failure modes
 207                            // that we know are safe to ignore here, like what we do
 208                            // with `CredentialsNotFound` above.
 209                            match provider_id.0.as_ref() {
 210                                "lmstudio" | "ollama" => {
 211                                    // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".
 212                                    //
 213                                    // These fail noisily, so we don't log them.
 214                                }
 215                                "copilot_chat" => {
 216                                    // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.
 217                                }
 218                                _ => {
 219                                    log::error!(
 220                                        "Failed to authenticate provider: {}: {err:#}",
 221                                        provider_name.0
 222                                    );
 223                                }
 224                            }
 225                        }
 226                    }
 227                }
 228            }
 229        })
 230    }
 231}
 232
 233pub struct NativeAgent {
 234    /// Session ID -> Session mapping
 235    sessions: HashMap<acp::SessionId, Session>,
 236    history: Entity<HistoryStore>,
 237    /// Shared project context for all threads
 238    project_context: Entity<ProjectContext>,
 239    project_context_needs_refresh: watch::Sender<()>,
 240    _maintain_project_context: Task<Result<()>>,
 241    context_server_registry: Entity<ContextServerRegistry>,
 242    /// Shared templates for all threads
 243    templates: Arc<Templates>,
 244    /// Cached model information
 245    models: LanguageModels,
 246    project: Entity<Project>,
 247    prompt_store: Option<Entity<PromptStore>>,
 248    fs: Arc<dyn Fs>,
 249    _subscriptions: Vec<Subscription>,
 250}
 251
 252impl NativeAgent {
 253    pub async fn new(
 254        project: Entity<Project>,
 255        history: Entity<HistoryStore>,
 256        templates: Arc<Templates>,
 257        prompt_store: Option<Entity<PromptStore>>,
 258        fs: Arc<dyn Fs>,
 259        cx: &mut AsyncApp,
 260    ) -> Result<Entity<NativeAgent>> {
 261        log::debug!("Creating new NativeAgent");
 262
 263        let project_context = cx
 264            .update(|cx| Self::build_project_context(&project, prompt_store.as_ref(), cx))?
 265            .await;
 266
 267        cx.new(|cx| {
 268            let mut subscriptions = vec![
 269                cx.subscribe(&project, Self::handle_project_event),
 270                cx.subscribe(
 271                    &LanguageModelRegistry::global(cx),
 272                    Self::handle_models_updated_event,
 273                ),
 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            let (project_context_needs_refresh_tx, project_context_needs_refresh_rx) =
 280                watch::channel(());
 281            Self {
 282                sessions: HashMap::new(),
 283                history,
 284                project_context: cx.new(|_| project_context),
 285                project_context_needs_refresh: project_context_needs_refresh_tx,
 286                _maintain_project_context: cx.spawn(async move |this, cx| {
 287                    Self::maintain_project_context(this, project_context_needs_refresh_rx, cx).await
 288                }),
 289                context_server_registry: cx.new(|cx| {
 290                    ContextServerRegistry::new(project.read(cx).context_server_store(), cx)
 291                }),
 292                templates,
 293                models: LanguageModels::new(cx),
 294                project,
 295                prompt_store,
 296                fs,
 297                _subscriptions: subscriptions,
 298            }
 299        })
 300    }
 301
 302    fn register_session(
 303        &mut self,
 304        thread_handle: Entity<Thread>,
 305        cx: &mut Context<Self>,
 306    ) -> Entity<AcpThread> {
 307        let connection = Rc::new(NativeAgentConnection(cx.entity()));
 308
 309        let thread = thread_handle.read(cx);
 310        let session_id = thread.id().clone();
 311        let title = thread.title();
 312        let project = thread.project.clone();
 313        let action_log = thread.action_log.clone();
 314        let prompt_capabilities_rx = thread.prompt_capabilities_rx.clone();
 315        let acp_thread = cx.new(|cx| {
 316            acp_thread::AcpThread::new(
 317                title,
 318                connection,
 319                project.clone(),
 320                action_log.clone(),
 321                session_id.clone(),
 322                prompt_capabilities_rx,
 323                cx,
 324            )
 325        });
 326
 327        let registry = LanguageModelRegistry::read_global(cx);
 328        let summarization_model = registry.thread_summary_model().map(|c| c.model);
 329
 330        thread_handle.update(cx, |thread, cx| {
 331            thread.set_summarization_model(summarization_model, cx);
 332            thread.add_default_tools(
 333                Rc::new(AcpThreadEnvironment {
 334                    acp_thread: acp_thread.downgrade(),
 335                }) as _,
 336                cx,
 337            )
 338        });
 339
 340        let subscriptions = vec![
 341            cx.observe_release(&acp_thread, |this, acp_thread, _cx| {
 342                this.sessions.remove(acp_thread.session_id());
 343            }),
 344            cx.subscribe(&thread_handle, Self::handle_thread_title_updated),
 345            cx.subscribe(&thread_handle, Self::handle_thread_token_usage_updated),
 346            cx.observe(&thread_handle, move |this, thread, cx| {
 347                this.save_thread(thread, cx)
 348            }),
 349        ];
 350
 351        self.sessions.insert(
 352            session_id,
 353            Session {
 354                thread: thread_handle,
 355                acp_thread: acp_thread.downgrade(),
 356                _subscriptions: subscriptions,
 357                pending_save: Task::ready(()),
 358            },
 359        );
 360        acp_thread
 361    }
 362
 363    pub fn models(&self) -> &LanguageModels {
 364        &self.models
 365    }
 366
 367    async fn maintain_project_context(
 368        this: WeakEntity<Self>,
 369        mut needs_refresh: watch::Receiver<()>,
 370        cx: &mut AsyncApp,
 371    ) -> Result<()> {
 372        while needs_refresh.changed().await.is_ok() {
 373            let project_context = this
 374                .update(cx, |this, cx| {
 375                    Self::build_project_context(&this.project, this.prompt_store.as_ref(), cx)
 376                })?
 377                .await;
 378            this.update(cx, |this, cx| {
 379                this.project_context = cx.new(|_| project_context);
 380            })?;
 381        }
 382
 383        Ok(())
 384    }
 385
 386    fn build_project_context(
 387        project: &Entity<Project>,
 388        prompt_store: Option<&Entity<PromptStore>>,
 389        cx: &mut App,
 390    ) -> Task<ProjectContext> {
 391        let worktrees = project.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
 392        let worktree_tasks = worktrees
 393            .into_iter()
 394            .map(|worktree| {
 395                Self::load_worktree_info_for_system_prompt(worktree, project.clone(), cx)
 396            })
 397            .collect::<Vec<_>>();
 398        let default_user_rules_task = if let Some(prompt_store) = prompt_store.as_ref() {
 399            prompt_store.read_with(cx, |prompt_store, cx| {
 400                let prompts = prompt_store.default_prompt_metadata();
 401                let load_tasks = prompts.into_iter().map(|prompt_metadata| {
 402                    let contents = prompt_store.load(prompt_metadata.id, cx);
 403                    async move { (contents.await, prompt_metadata) }
 404                });
 405                cx.background_spawn(future::join_all(load_tasks))
 406            })
 407        } else {
 408            Task::ready(vec![])
 409        };
 410
 411        cx.spawn(async move |_cx| {
 412            let (worktrees, default_user_rules) =
 413                future::join(future::join_all(worktree_tasks), default_user_rules_task).await;
 414
 415            let worktrees = worktrees
 416                .into_iter()
 417                .map(|(worktree, _rules_error)| {
 418                    // TODO: show error message
 419                    // if let Some(rules_error) = rules_error {
 420                    //     this.update(cx, |_, cx| cx.emit(rules_error)).ok();
 421                    // }
 422                    worktree
 423                })
 424                .collect::<Vec<_>>();
 425
 426            let default_user_rules = default_user_rules
 427                .into_iter()
 428                .flat_map(|(contents, prompt_metadata)| match contents {
 429                    Ok(contents) => Some(UserRulesContext {
 430                        uuid: match prompt_metadata.id {
 431                            prompt_store::PromptId::User { uuid } => uuid,
 432                            prompt_store::PromptId::EditWorkflow => return None,
 433                        },
 434                        title: prompt_metadata.title.map(|title| title.to_string()),
 435                        contents,
 436                    }),
 437                    Err(_err) => {
 438                        // TODO: show error message
 439                        // this.update(cx, |_, cx| {
 440                        //     cx.emit(RulesLoadingError {
 441                        //         message: format!("{err:?}").into(),
 442                        //     });
 443                        // })
 444                        // .ok();
 445                        None
 446                    }
 447                })
 448                .collect::<Vec<_>>();
 449
 450            ProjectContext::new(worktrees, default_user_rules)
 451        })
 452    }
 453
 454    fn load_worktree_info_for_system_prompt(
 455        worktree: Entity<Worktree>,
 456        project: Entity<Project>,
 457        cx: &mut App,
 458    ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
 459        let tree = worktree.read(cx);
 460        let root_name = tree.root_name_str().into();
 461        let abs_path = tree.abs_path();
 462
 463        let mut context = WorktreeContext {
 464            root_name,
 465            abs_path,
 466            rules_file: None,
 467        };
 468
 469        let rules_task = Self::load_worktree_rules_file(worktree, project, cx);
 470        let Some(rules_task) = rules_task else {
 471            return Task::ready((context, None));
 472        };
 473
 474        cx.spawn(async move |_| {
 475            let (rules_file, rules_file_error) = match rules_task.await {
 476                Ok(rules_file) => (Some(rules_file), None),
 477                Err(err) => (
 478                    None,
 479                    Some(RulesLoadingError {
 480                        message: format!("{err}").into(),
 481                    }),
 482                ),
 483            };
 484            context.rules_file = rules_file;
 485            (context, rules_file_error)
 486        })
 487    }
 488
 489    fn load_worktree_rules_file(
 490        worktree: Entity<Worktree>,
 491        project: Entity<Project>,
 492        cx: &mut App,
 493    ) -> Option<Task<Result<RulesFileContext>>> {
 494        let worktree = worktree.read(cx);
 495        let worktree_id = worktree.id();
 496        let selected_rules_file = RULES_FILE_NAMES
 497            .into_iter()
 498            .filter_map(|name| {
 499                worktree
 500                    .entry_for_path(RelPath::unix(name).unwrap())
 501                    .filter(|entry| entry.is_file())
 502                    .map(|entry| entry.path.clone())
 503            })
 504            .next();
 505
 506        // Note that Cline supports `.clinerules` being a directory, but that is not currently
 507        // supported. This doesn't seem to occur often in GitHub repositories.
 508        selected_rules_file.map(|path_in_worktree| {
 509            let project_path = ProjectPath {
 510                worktree_id,
 511                path: path_in_worktree.clone(),
 512            };
 513            let buffer_task =
 514                project.update(cx, |project, cx| project.open_buffer(project_path, cx));
 515            let rope_task = cx.spawn(async move |cx| {
 516                buffer_task.await?.read_with(cx, |buffer, cx| {
 517                    let project_entry_id = buffer.entry_id(cx).context("buffer has no file")?;
 518                    anyhow::Ok((project_entry_id, buffer.as_rope().clone()))
 519                })?
 520            });
 521            // Build a string from the rope on a background thread.
 522            cx.background_spawn(async move {
 523                let (project_entry_id, rope) = rope_task.await?;
 524                anyhow::Ok(RulesFileContext {
 525                    path_in_worktree,
 526                    text: rope.to_string().trim().to_string(),
 527                    project_entry_id: project_entry_id.to_usize(),
 528                })
 529            })
 530        })
 531    }
 532
 533    fn handle_thread_title_updated(
 534        &mut self,
 535        thread: Entity<Thread>,
 536        _: &TitleUpdated,
 537        cx: &mut Context<Self>,
 538    ) {
 539        let session_id = thread.read(cx).id();
 540        let Some(session) = self.sessions.get(session_id) else {
 541            return;
 542        };
 543        let thread = thread.downgrade();
 544        let acp_thread = session.acp_thread.clone();
 545        cx.spawn(async move |_, cx| {
 546            let title = thread.read_with(cx, |thread, _| thread.title())?;
 547            let task = acp_thread.update(cx, |acp_thread, cx| acp_thread.set_title(title, cx))?;
 548            task.await
 549        })
 550        .detach_and_log_err(cx);
 551    }
 552
 553    fn handle_thread_token_usage_updated(
 554        &mut self,
 555        thread: Entity<Thread>,
 556        usage: &TokenUsageUpdated,
 557        cx: &mut Context<Self>,
 558    ) {
 559        let Some(session) = self.sessions.get(thread.read(cx).id()) else {
 560            return;
 561        };
 562        session
 563            .acp_thread
 564            .update(cx, |acp_thread, cx| {
 565                acp_thread.update_token_usage(usage.0.clone(), cx);
 566            })
 567            .ok();
 568    }
 569
 570    fn handle_project_event(
 571        &mut self,
 572        _project: Entity<Project>,
 573        event: &project::Event,
 574        _cx: &mut Context<Self>,
 575    ) {
 576        match event {
 577            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
 578                self.project_context_needs_refresh.send(()).ok();
 579            }
 580            project::Event::WorktreeUpdatedEntries(_, items) => {
 581                if items.iter().any(|(path, _, _)| {
 582                    RULES_FILE_NAMES
 583                        .iter()
 584                        .any(|name| path.as_ref() == RelPath::unix(name).unwrap())
 585                }) {
 586                    self.project_context_needs_refresh.send(()).ok();
 587                }
 588            }
 589            _ => {}
 590        }
 591    }
 592
 593    fn handle_prompts_updated_event(
 594        &mut self,
 595        _prompt_store: Entity<PromptStore>,
 596        _event: &prompt_store::PromptsUpdatedEvent,
 597        _cx: &mut Context<Self>,
 598    ) {
 599        self.project_context_needs_refresh.send(()).ok();
 600    }
 601
 602    fn handle_models_updated_event(
 603        &mut self,
 604        _registry: Entity<LanguageModelRegistry>,
 605        _event: &language_model::Event,
 606        cx: &mut Context<Self>,
 607    ) {
 608        self.models.refresh_list(cx);
 609
 610        let registry = LanguageModelRegistry::read_global(cx);
 611        let default_model = registry.default_model().map(|m| m.model);
 612        let summarization_model = registry.thread_summary_model().map(|m| m.model);
 613
 614        for session in self.sessions.values_mut() {
 615            session.thread.update(cx, |thread, cx| {
 616                if thread.model().is_none()
 617                    && let Some(model) = default_model.clone()
 618                {
 619                    thread.set_model(model, cx);
 620                    cx.notify();
 621                }
 622                thread.set_summarization_model(summarization_model.clone(), cx);
 623            });
 624        }
 625    }
 626
 627    pub fn load_thread(
 628        &mut self,
 629        id: acp::SessionId,
 630        cx: &mut Context<Self>,
 631    ) -> Task<Result<Entity<Thread>>> {
 632        let database_future = ThreadsDatabase::connect(cx);
 633        cx.spawn(async move |this, cx| {
 634            let database = database_future.await.map_err(|err| anyhow!(err))?;
 635            let db_thread = database
 636                .load_thread(id.clone())
 637                .await?
 638                .with_context(|| format!("no thread found with ID: {id:?}"))?;
 639
 640            this.update(cx, |this, cx| {
 641                let summarization_model = LanguageModelRegistry::read_global(cx)
 642                    .thread_summary_model()
 643                    .map(|c| c.model);
 644
 645                cx.new(|cx| {
 646                    let mut thread = Thread::from_db(
 647                        id.clone(),
 648                        db_thread,
 649                        this.project.clone(),
 650                        this.project_context.clone(),
 651                        this.context_server_registry.clone(),
 652                        this.templates.clone(),
 653                        cx,
 654                    );
 655                    thread.set_summarization_model(summarization_model, cx);
 656                    thread
 657                })
 658            })
 659        })
 660    }
 661
 662    pub fn open_thread(
 663        &mut self,
 664        id: acp::SessionId,
 665        cx: &mut Context<Self>,
 666    ) -> Task<Result<Entity<AcpThread>>> {
 667        let task = self.load_thread(id, cx);
 668        cx.spawn(async move |this, cx| {
 669            let thread = task.await?;
 670            let acp_thread =
 671                this.update(cx, |this, cx| this.register_session(thread.clone(), cx))?;
 672            let events = thread.update(cx, |thread, cx| thread.replay(cx))?;
 673            cx.update(|cx| {
 674                NativeAgentConnection::handle_thread_events(events, acp_thread.downgrade(), cx)
 675            })?
 676            .await?;
 677            Ok(acp_thread)
 678        })
 679    }
 680
 681    pub fn thread_summary(
 682        &mut self,
 683        id: acp::SessionId,
 684        cx: &mut Context<Self>,
 685    ) -> Task<Result<SharedString>> {
 686        let thread = self.open_thread(id.clone(), cx);
 687        cx.spawn(async move |this, cx| {
 688            let acp_thread = thread.await?;
 689            let result = this
 690                .update(cx, |this, cx| {
 691                    this.sessions
 692                        .get(&id)
 693                        .unwrap()
 694                        .thread
 695                        .update(cx, |thread, cx| thread.summary(cx))
 696                })?
 697                .await
 698                .context("Failed to generate summary")?;
 699            drop(acp_thread);
 700            Ok(result)
 701        })
 702    }
 703
 704    fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
 705        if thread.read(cx).is_empty() {
 706            return;
 707        }
 708
 709        let database_future = ThreadsDatabase::connect(cx);
 710        let (id, db_thread) =
 711            thread.update(cx, |thread, cx| (thread.id().clone(), thread.to_db(cx)));
 712        let Some(session) = self.sessions.get_mut(&id) else {
 713            return;
 714        };
 715        let history = self.history.clone();
 716        session.pending_save = cx.spawn(async move |_, cx| {
 717            let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
 718                return;
 719            };
 720            let db_thread = db_thread.await;
 721            database.save_thread(id, db_thread).await.log_err();
 722            history.update(cx, |history, cx| history.reload(cx)).ok();
 723        });
 724    }
 725}
 726
 727/// Wrapper struct that implements the AgentConnection trait
 728#[derive(Clone)]
 729pub struct NativeAgentConnection(pub Entity<NativeAgent>);
 730
 731impl NativeAgentConnection {
 732    pub fn thread(&self, session_id: &acp::SessionId, cx: &App) -> Option<Entity<Thread>> {
 733        self.0
 734            .read(cx)
 735            .sessions
 736            .get(session_id)
 737            .map(|session| session.thread.clone())
 738    }
 739
 740    pub fn load_thread(&self, id: acp::SessionId, cx: &mut App) -> Task<Result<Entity<Thread>>> {
 741        self.0.update(cx, |this, cx| this.load_thread(id, cx))
 742    }
 743
 744    fn run_turn(
 745        &self,
 746        session_id: acp::SessionId,
 747        cx: &mut App,
 748        f: impl 'static
 749        + FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
 750    ) -> Task<Result<acp::PromptResponse>> {
 751        let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
 752            agent
 753                .sessions
 754                .get_mut(&session_id)
 755                .map(|s| (s.thread.clone(), s.acp_thread.clone()))
 756        }) else {
 757            return Task::ready(Err(anyhow!("Session not found")));
 758        };
 759        log::debug!("Found session for: {}", session_id);
 760
 761        let response_stream = match f(thread, cx) {
 762            Ok(stream) => stream,
 763            Err(err) => return Task::ready(Err(err)),
 764        };
 765        Self::handle_thread_events(response_stream, acp_thread, cx)
 766    }
 767
 768    fn handle_thread_events(
 769        mut events: mpsc::UnboundedReceiver<Result<ThreadEvent>>,
 770        acp_thread: WeakEntity<AcpThread>,
 771        cx: &App,
 772    ) -> Task<Result<acp::PromptResponse>> {
 773        cx.spawn(async move |cx| {
 774            // Handle response stream and forward to session.acp_thread
 775            while let Some(result) = events.next().await {
 776                match result {
 777                    Ok(event) => {
 778                        log::trace!("Received completion event: {:?}", event);
 779
 780                        match event {
 781                            ThreadEvent::UserMessage(message) => {
 782                                acp_thread.update(cx, |thread, cx| {
 783                                    for content in message.content {
 784                                        thread.push_user_content_block(
 785                                            Some(message.id.clone()),
 786                                            content.into(),
 787                                            cx,
 788                                        );
 789                                    }
 790                                })?;
 791                            }
 792                            ThreadEvent::AgentText(text) => {
 793                                acp_thread.update(cx, |thread, cx| {
 794                                    thread.push_assistant_content_block(
 795                                        acp::ContentBlock::Text(acp::TextContent {
 796                                            text,
 797                                            annotations: None,
 798                                            meta: None,
 799                                        }),
 800                                        false,
 801                                        cx,
 802                                    )
 803                                })?;
 804                            }
 805                            ThreadEvent::AgentThinking(text) => {
 806                                acp_thread.update(cx, |thread, cx| {
 807                                    thread.push_assistant_content_block(
 808                                        acp::ContentBlock::Text(acp::TextContent {
 809                                            text,
 810                                            annotations: None,
 811                                            meta: None,
 812                                        }),
 813                                        true,
 814                                        cx,
 815                                    )
 816                                })?;
 817                            }
 818                            ThreadEvent::ToolCallAuthorization(ToolCallAuthorization {
 819                                tool_call,
 820                                options,
 821                                response,
 822                            }) => {
 823                                let outcome_task = acp_thread.update(cx, |thread, cx| {
 824                                    thread.request_tool_call_authorization(
 825                                        tool_call, options, true, cx,
 826                                    )
 827                                })??;
 828                                cx.background_spawn(async move {
 829                                    if let acp::RequestPermissionOutcome::Selected { option_id } =
 830                                        outcome_task.await
 831                                    {
 832                                        response
 833                                            .send(option_id)
 834                                            .map(|_| anyhow!("authorization receiver was dropped"))
 835                                            .log_err();
 836                                    }
 837                                })
 838                                .detach();
 839                            }
 840                            ThreadEvent::ToolCall(tool_call) => {
 841                                acp_thread.update(cx, |thread, cx| {
 842                                    thread.upsert_tool_call(tool_call, cx)
 843                                })??;
 844                            }
 845                            ThreadEvent::ToolCallUpdate(update) => {
 846                                acp_thread.update(cx, |thread, cx| {
 847                                    thread.update_tool_call(update, cx)
 848                                })??;
 849                            }
 850                            ThreadEvent::Retry(status) => {
 851                                acp_thread.update(cx, |thread, cx| {
 852                                    thread.update_retry_status(status, cx)
 853                                })?;
 854                            }
 855                            ThreadEvent::Stop(stop_reason) => {
 856                                log::debug!("Assistant message complete: {:?}", stop_reason);
 857                                return Ok(acp::PromptResponse {
 858                                    stop_reason,
 859                                    meta: None,
 860                                });
 861                            }
 862                        }
 863                    }
 864                    Err(e) => {
 865                        log::error!("Error in model response stream: {:?}", e);
 866                        return Err(e);
 867                    }
 868                }
 869            }
 870
 871            log::debug!("Response stream completed");
 872            anyhow::Ok(acp::PromptResponse {
 873                stop_reason: acp::StopReason::EndTurn,
 874                meta: None,
 875            })
 876        })
 877    }
 878}
 879
 880struct NativeAgentModelSelector {
 881    session_id: acp::SessionId,
 882    connection: NativeAgentConnection,
 883}
 884
 885impl acp_thread::AgentModelSelector for NativeAgentModelSelector {
 886    fn list_models(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelList>> {
 887        log::debug!("NativeAgentConnection::list_models called");
 888        let list = self.connection.0.read(cx).models.model_list.clone();
 889        Task::ready(if list.is_empty() {
 890            Err(anyhow::anyhow!("No models available"))
 891        } else {
 892            Ok(list)
 893        })
 894    }
 895
 896    fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>> {
 897        log::debug!(
 898            "Setting model for session {}: {}",
 899            self.session_id,
 900            model_id
 901        );
 902        let Some(thread) = self
 903            .connection
 904            .0
 905            .read(cx)
 906            .sessions
 907            .get(&self.session_id)
 908            .map(|session| session.thread.clone())
 909        else {
 910            return Task::ready(Err(anyhow!("Session not found")));
 911        };
 912
 913        let Some(model) = self.connection.0.read(cx).models.model_from_id(&model_id) else {
 914            return Task::ready(Err(anyhow!("Invalid model ID {}", model_id)));
 915        };
 916
 917        thread.update(cx, |thread, cx| {
 918            thread.set_model(model.clone(), cx);
 919        });
 920
 921        update_settings_file(
 922            self.connection.0.read(cx).fs.clone(),
 923            cx,
 924            move |settings, _cx| {
 925                let provider = model.provider_id().0.to_string();
 926                let model = model.id().0.to_string();
 927                settings
 928                    .agent
 929                    .get_or_insert_default()
 930                    .set_model(LanguageModelSelection {
 931                        provider: provider.into(),
 932                        model,
 933                    });
 934            },
 935        );
 936
 937        Task::ready(Ok(()))
 938    }
 939
 940    fn selected_model(&self, cx: &mut App) -> Task<Result<acp_thread::AgentModelInfo>> {
 941        let Some(thread) = self
 942            .connection
 943            .0
 944            .read(cx)
 945            .sessions
 946            .get(&self.session_id)
 947            .map(|session| session.thread.clone())
 948        else {
 949            return Task::ready(Err(anyhow!("Session not found")));
 950        };
 951        let Some(model) = thread.read(cx).model() else {
 952            return Task::ready(Err(anyhow!("Model not found")));
 953        };
 954        let Some(provider) = LanguageModelRegistry::read_global(cx).provider(&model.provider_id())
 955        else {
 956            return Task::ready(Err(anyhow!("Provider not found")));
 957        };
 958        Task::ready(Ok(LanguageModels::map_language_model_to_info(
 959            model, &provider,
 960        )))
 961    }
 962
 963    fn watch(&self, cx: &mut App) -> Option<watch::Receiver<()>> {
 964        Some(self.connection.0.read(cx).models.watch())
 965    }
 966}
 967
 968impl acp_thread::AgentConnection for NativeAgentConnection {
 969    fn telemetry_id(&self) -> &'static str {
 970        "zed"
 971    }
 972
 973    fn new_thread(
 974        self: Rc<Self>,
 975        project: Entity<Project>,
 976        cwd: &Path,
 977        cx: &mut App,
 978    ) -> Task<Result<Entity<acp_thread::AcpThread>>> {
 979        let agent = self.0.clone();
 980        log::debug!("Creating new thread for project at: {:?}", cwd);
 981
 982        cx.spawn(async move |cx| {
 983            log::debug!("Starting thread creation in async context");
 984
 985            // Create Thread
 986            let thread = agent.update(
 987                cx,
 988                |agent, cx: &mut gpui::Context<NativeAgent>| -> Result<_> {
 989                    // Fetch default model from registry settings
 990                    let registry = LanguageModelRegistry::read_global(cx);
 991                    // Log available models for debugging
 992                    let available_count = registry.available_models(cx).count();
 993                    log::debug!("Total available models: {}", available_count);
 994
 995                    let default_model = registry.default_model().and_then(|default_model| {
 996                        agent
 997                            .models
 998                            .model_from_id(&LanguageModels::model_id(&default_model.model))
 999                    });
1000                    Ok(cx.new(|cx| {
1001                        Thread::new(
1002                            project.clone(),
1003                            agent.project_context.clone(),
1004                            agent.context_server_registry.clone(),
1005                            agent.templates.clone(),
1006                            default_model,
1007                            cx,
1008                        )
1009                    }))
1010                },
1011            )??;
1012            agent.update(cx, |agent, cx| agent.register_session(thread, cx))
1013        })
1014    }
1015
1016    fn auth_methods(&self) -> &[acp::AuthMethod] {
1017        &[] // No auth for in-process
1018    }
1019
1020    fn authenticate(&self, _method: acp::AuthMethodId, _cx: &mut App) -> Task<Result<()>> {
1021        Task::ready(Ok(()))
1022    }
1023
1024    fn model_selector(&self, session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
1025        Some(Rc::new(NativeAgentModelSelector {
1026            session_id: session_id.clone(),
1027            connection: self.clone(),
1028        }) as Rc<dyn AgentModelSelector>)
1029    }
1030
1031    fn prompt(
1032        &self,
1033        id: Option<acp_thread::UserMessageId>,
1034        params: acp::PromptRequest,
1035        cx: &mut App,
1036    ) -> Task<Result<acp::PromptResponse>> {
1037        let id = id.expect("UserMessageId is required");
1038        let session_id = params.session_id.clone();
1039        log::info!("Received prompt request for session: {}", session_id);
1040        log::debug!("Prompt blocks count: {}", params.prompt.len());
1041        let path_style = self.0.read(cx).project.read(cx).path_style(cx);
1042
1043        self.run_turn(session_id, cx, move |thread, cx| {
1044            let content: Vec<UserMessageContent> = params
1045                .prompt
1046                .into_iter()
1047                .map(|block| UserMessageContent::from_content_block(block, path_style))
1048                .collect::<Vec<_>>();
1049            log::debug!("Converted prompt to message: {} chars", content.len());
1050            log::debug!("Message id: {:?}", id);
1051            log::debug!("Message content: {:?}", content);
1052
1053            thread.update(cx, |thread, cx| thread.send(id, content, cx))
1054        })
1055    }
1056
1057    fn resume(
1058        &self,
1059        session_id: &acp::SessionId,
1060        _cx: &App,
1061    ) -> Option<Rc<dyn acp_thread::AgentSessionResume>> {
1062        Some(Rc::new(NativeAgentSessionResume {
1063            connection: self.clone(),
1064            session_id: session_id.clone(),
1065        }) as _)
1066    }
1067
1068    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
1069        log::info!("Cancelling on session: {}", session_id);
1070        self.0.update(cx, |agent, cx| {
1071            if let Some(agent) = agent.sessions.get(session_id) {
1072                agent.thread.update(cx, |thread, cx| thread.cancel(cx));
1073            }
1074        });
1075    }
1076
1077    fn truncate(
1078        &self,
1079        session_id: &agent_client_protocol::SessionId,
1080        cx: &App,
1081    ) -> Option<Rc<dyn acp_thread::AgentSessionTruncate>> {
1082        self.0.read_with(cx, |agent, _cx| {
1083            agent.sessions.get(session_id).map(|session| {
1084                Rc::new(NativeAgentSessionTruncate {
1085                    thread: session.thread.clone(),
1086                    acp_thread: session.acp_thread.clone(),
1087                }) as _
1088            })
1089        })
1090    }
1091
1092    fn set_title(
1093        &self,
1094        session_id: &acp::SessionId,
1095        _cx: &App,
1096    ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
1097        Some(Rc::new(NativeAgentSessionSetTitle {
1098            connection: self.clone(),
1099            session_id: session_id.clone(),
1100        }) as _)
1101    }
1102
1103    fn telemetry(&self) -> Option<Rc<dyn acp_thread::AgentTelemetry>> {
1104        Some(Rc::new(self.clone()) as Rc<dyn acp_thread::AgentTelemetry>)
1105    }
1106
1107    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1108        self
1109    }
1110}
1111
1112impl acp_thread::AgentTelemetry for NativeAgentConnection {
1113    fn thread_data(
1114        &self,
1115        session_id: &acp::SessionId,
1116        cx: &mut App,
1117    ) -> Task<Result<serde_json::Value>> {
1118        let Some(session) = self.0.read(cx).sessions.get(session_id) else {
1119            return Task::ready(Err(anyhow!("Session not found")));
1120        };
1121
1122        let task = session.thread.read(cx).to_db(cx);
1123        cx.background_spawn(async move {
1124            serde_json::to_value(task.await).context("Failed to serialize thread")
1125        })
1126    }
1127}
1128
1129struct NativeAgentSessionTruncate {
1130    thread: Entity<Thread>,
1131    acp_thread: WeakEntity<AcpThread>,
1132}
1133
1134impl acp_thread::AgentSessionTruncate for NativeAgentSessionTruncate {
1135    fn run(&self, message_id: acp_thread::UserMessageId, cx: &mut App) -> Task<Result<()>> {
1136        match self.thread.update(cx, |thread, cx| {
1137            thread.truncate(message_id.clone(), cx)?;
1138            Ok(thread.latest_token_usage())
1139        }) {
1140            Ok(usage) => {
1141                self.acp_thread
1142                    .update(cx, |thread, cx| {
1143                        thread.update_token_usage(usage, cx);
1144                    })
1145                    .ok();
1146                Task::ready(Ok(()))
1147            }
1148            Err(error) => Task::ready(Err(error)),
1149        }
1150    }
1151}
1152
1153struct NativeAgentSessionResume {
1154    connection: NativeAgentConnection,
1155    session_id: acp::SessionId,
1156}
1157
1158impl acp_thread::AgentSessionResume for NativeAgentSessionResume {
1159    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>> {
1160        self.connection
1161            .run_turn(self.session_id.clone(), cx, |thread, cx| {
1162                thread.update(cx, |thread, cx| thread.resume(cx))
1163            })
1164    }
1165}
1166
1167struct NativeAgentSessionSetTitle {
1168    connection: NativeAgentConnection,
1169    session_id: acp::SessionId,
1170}
1171
1172impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
1173    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
1174        let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
1175            return Task::ready(Err(anyhow!("session not found")));
1176        };
1177        let thread = session.thread.clone();
1178        thread.update(cx, |thread, cx| thread.set_title(title, cx));
1179        Task::ready(Ok(()))
1180    }
1181}
1182
1183pub struct AcpThreadEnvironment {
1184    acp_thread: WeakEntity<AcpThread>,
1185}
1186
1187impl ThreadEnvironment for AcpThreadEnvironment {
1188    fn create_terminal(
1189        &self,
1190        command: String,
1191        cwd: Option<PathBuf>,
1192        output_byte_limit: Option<u64>,
1193        cx: &mut AsyncApp,
1194    ) -> Task<Result<Rc<dyn TerminalHandle>>> {
1195        let task = self.acp_thread.update(cx, |thread, cx| {
1196            thread.create_terminal(command, vec![], vec![], cwd, output_byte_limit, cx)
1197        });
1198
1199        let acp_thread = self.acp_thread.clone();
1200        cx.spawn(async move |cx| {
1201            let terminal = task?.await?;
1202
1203            let (drop_tx, drop_rx) = oneshot::channel();
1204            let terminal_id = terminal.read_with(cx, |terminal, _cx| terminal.id().clone())?;
1205
1206            cx.spawn(async move |cx| {
1207                drop_rx.await.ok();
1208                acp_thread.update(cx, |thread, cx| thread.release_terminal(terminal_id, cx))
1209            })
1210            .detach();
1211
1212            let handle = AcpTerminalHandle {
1213                terminal,
1214                _drop_tx: Some(drop_tx),
1215            };
1216
1217            Ok(Rc::new(handle) as _)
1218        })
1219    }
1220}
1221
1222pub struct AcpTerminalHandle {
1223    terminal: Entity<acp_thread::Terminal>,
1224    _drop_tx: Option<oneshot::Sender<()>>,
1225}
1226
1227impl TerminalHandle for AcpTerminalHandle {
1228    fn id(&self, cx: &AsyncApp) -> Result<acp::TerminalId> {
1229        self.terminal.read_with(cx, |term, _cx| term.id().clone())
1230    }
1231
1232    fn wait_for_exit(&self, cx: &AsyncApp) -> Result<Shared<Task<acp::TerminalExitStatus>>> {
1233        self.terminal
1234            .read_with(cx, |term, _cx| term.wait_for_exit())
1235    }
1236
1237    fn current_output(&self, cx: &AsyncApp) -> Result<acp::TerminalOutputResponse> {
1238        self.terminal
1239            .read_with(cx, |term, cx| term.current_output(cx))
1240    }
1241}
1242
1243#[cfg(test)]
1244mod internal_tests {
1245    use crate::HistoryEntryId;
1246
1247    use super::*;
1248    use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
1249    use fs::FakeFs;
1250    use gpui::TestAppContext;
1251    use indoc::formatdoc;
1252    use language_model::fake_provider::FakeLanguageModel;
1253    use serde_json::json;
1254    use settings::SettingsStore;
1255    use util::{path, rel_path::rel_path};
1256
1257    #[gpui::test]
1258    async fn test_maintaining_project_context(cx: &mut TestAppContext) {
1259        init_test(cx);
1260        let fs = FakeFs::new(cx.executor());
1261        fs.insert_tree(
1262            "/",
1263            json!({
1264                "a": {}
1265            }),
1266        )
1267        .await;
1268        let project = Project::test(fs.clone(), [], cx).await;
1269        let text_thread_store =
1270            cx.new(|cx| assistant_text_thread::TextThreadStore::fake(project.clone(), cx));
1271        let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
1272        let agent = NativeAgent::new(
1273            project.clone(),
1274            history_store,
1275            Templates::new(),
1276            None,
1277            fs.clone(),
1278            &mut cx.to_async(),
1279        )
1280        .await
1281        .unwrap();
1282        agent.read_with(cx, |agent, cx| {
1283            assert_eq!(agent.project_context.read(cx).worktrees, vec![])
1284        });
1285
1286        let worktree = project
1287            .update(cx, |project, cx| project.create_worktree("/a", true, cx))
1288            .await
1289            .unwrap();
1290        cx.run_until_parked();
1291        agent.read_with(cx, |agent, cx| {
1292            assert_eq!(
1293                agent.project_context.read(cx).worktrees,
1294                vec![WorktreeContext {
1295                    root_name: "a".into(),
1296                    abs_path: Path::new("/a").into(),
1297                    rules_file: None
1298                }]
1299            )
1300        });
1301
1302        // Creating `/a/.rules` updates the project context.
1303        fs.insert_file("/a/.rules", Vec::new()).await;
1304        cx.run_until_parked();
1305        agent.read_with(cx, |agent, cx| {
1306            let rules_entry = worktree
1307                .read(cx)
1308                .entry_for_path(rel_path(".rules"))
1309                .unwrap();
1310            assert_eq!(
1311                agent.project_context.read(cx).worktrees,
1312                vec![WorktreeContext {
1313                    root_name: "a".into(),
1314                    abs_path: Path::new("/a").into(),
1315                    rules_file: Some(RulesFileContext {
1316                        path_in_worktree: rel_path(".rules").into(),
1317                        text: "".into(),
1318                        project_entry_id: rules_entry.id.to_usize()
1319                    })
1320                }]
1321            )
1322        });
1323    }
1324
1325    #[gpui::test]
1326    async fn test_listing_models(cx: &mut TestAppContext) {
1327        init_test(cx);
1328        let fs = FakeFs::new(cx.executor());
1329        fs.insert_tree("/", json!({ "a": {}  })).await;
1330        let project = Project::test(fs.clone(), [], cx).await;
1331        let text_thread_store =
1332            cx.new(|cx| assistant_text_thread::TextThreadStore::fake(project.clone(), cx));
1333        let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
1334        let connection = NativeAgentConnection(
1335            NativeAgent::new(
1336                project.clone(),
1337                history_store,
1338                Templates::new(),
1339                None,
1340                fs.clone(),
1341                &mut cx.to_async(),
1342            )
1343            .await
1344            .unwrap(),
1345        );
1346
1347        // Create a thread/session
1348        let acp_thread = cx
1349            .update(|cx| {
1350                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1351            })
1352            .await
1353            .unwrap();
1354
1355        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1356
1357        let models = cx
1358            .update(|cx| {
1359                connection
1360                    .model_selector(&session_id)
1361                    .unwrap()
1362                    .list_models(cx)
1363            })
1364            .await
1365            .unwrap();
1366
1367        let acp_thread::AgentModelList::Grouped(models) = models else {
1368            panic!("Unexpected model group");
1369        };
1370        assert_eq!(
1371            models,
1372            IndexMap::from_iter([(
1373                AgentModelGroupName("Fake".into()),
1374                vec![AgentModelInfo {
1375                    id: acp::ModelId("fake/fake".into()),
1376                    name: "Fake".into(),
1377                    description: None,
1378                    icon: Some(ui::IconName::ZedAssistant),
1379                }]
1380            )])
1381        );
1382    }
1383
1384    #[gpui::test]
1385    async fn test_model_selection_persists_to_settings(cx: &mut TestAppContext) {
1386        init_test(cx);
1387        let fs = FakeFs::new(cx.executor());
1388        fs.create_dir(paths::settings_file().parent().unwrap())
1389            .await
1390            .unwrap();
1391        fs.insert_file(
1392            paths::settings_file(),
1393            json!({
1394                "agent": {
1395                    "default_model": {
1396                        "provider": "foo",
1397                        "model": "bar"
1398                    }
1399                }
1400            })
1401            .to_string()
1402            .into_bytes(),
1403        )
1404        .await;
1405        let project = Project::test(fs.clone(), [], cx).await;
1406
1407        let text_thread_store =
1408            cx.new(|cx| assistant_text_thread::TextThreadStore::fake(project.clone(), cx));
1409        let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
1410
1411        // Create the agent and connection
1412        let agent = NativeAgent::new(
1413            project.clone(),
1414            history_store,
1415            Templates::new(),
1416            None,
1417            fs.clone(),
1418            &mut cx.to_async(),
1419        )
1420        .await
1421        .unwrap();
1422        let connection = NativeAgentConnection(agent.clone());
1423
1424        // Create a thread/session
1425        let acp_thread = cx
1426            .update(|cx| {
1427                Rc::new(connection.clone()).new_thread(project.clone(), Path::new("/a"), cx)
1428            })
1429            .await
1430            .unwrap();
1431
1432        let session_id = cx.update(|cx| acp_thread.read(cx).session_id().clone());
1433
1434        // Select a model
1435        let selector = connection.model_selector(&session_id).unwrap();
1436        let model_id = acp::ModelId("fake/fake".into());
1437        cx.update(|cx| selector.select_model(model_id.clone(), cx))
1438            .await
1439            .unwrap();
1440
1441        // Verify the thread has the selected model
1442        agent.read_with(cx, |agent, _| {
1443            let session = agent.sessions.get(&session_id).unwrap();
1444            session.thread.read_with(cx, |thread, _| {
1445                assert_eq!(thread.model().unwrap().id().0, "fake");
1446            });
1447        });
1448
1449        cx.run_until_parked();
1450
1451        // Verify settings file was updated
1452        let settings_content = fs.load(paths::settings_file()).await.unwrap();
1453        let settings_json: serde_json::Value = serde_json::from_str(&settings_content).unwrap();
1454
1455        // Check that the agent settings contain the selected model
1456        assert_eq!(
1457            settings_json["agent"]["default_model"]["model"],
1458            json!("fake")
1459        );
1460        assert_eq!(
1461            settings_json["agent"]["default_model"]["provider"],
1462            json!("fake")
1463        );
1464    }
1465
1466    #[gpui::test]
1467    async fn test_save_load_thread(cx: &mut TestAppContext) {
1468        init_test(cx);
1469        let fs = FakeFs::new(cx.executor());
1470        fs.insert_tree(
1471            "/",
1472            json!({
1473                "a": {
1474                    "b.md": "Lorem"
1475                }
1476            }),
1477        )
1478        .await;
1479        let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
1480        let text_thread_store =
1481            cx.new(|cx| assistant_text_thread::TextThreadStore::fake(project.clone(), cx));
1482        let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
1483        let agent = NativeAgent::new(
1484            project.clone(),
1485            history_store.clone(),
1486            Templates::new(),
1487            None,
1488            fs.clone(),
1489            &mut cx.to_async(),
1490        )
1491        .await
1492        .unwrap();
1493        let connection = Rc::new(NativeAgentConnection(agent.clone()));
1494
1495        let acp_thread = cx
1496            .update(|cx| {
1497                connection
1498                    .clone()
1499                    .new_thread(project.clone(), Path::new(""), cx)
1500            })
1501            .await
1502            .unwrap();
1503        let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
1504        let thread = agent.read_with(cx, |agent, _| {
1505            agent.sessions.get(&session_id).unwrap().thread.clone()
1506        });
1507
1508        // Ensure empty threads are not saved, even if they get mutated.
1509        let model = Arc::new(FakeLanguageModel::default());
1510        let summary_model = Arc::new(FakeLanguageModel::default());
1511        thread.update(cx, |thread, cx| {
1512            thread.set_model(model.clone(), cx);
1513            thread.set_summarization_model(Some(summary_model.clone()), cx);
1514        });
1515        cx.run_until_parked();
1516        assert_eq!(history_entries(&history_store, cx), vec![]);
1517
1518        let send = acp_thread.update(cx, |thread, cx| {
1519            thread.send(
1520                vec![
1521                    "What does ".into(),
1522                    acp::ContentBlock::ResourceLink(acp::ResourceLink {
1523                        name: "b.md".into(),
1524                        uri: MentionUri::File {
1525                            abs_path: path!("/a/b.md").into(),
1526                        }
1527                        .to_uri()
1528                        .to_string(),
1529                        annotations: None,
1530                        description: None,
1531                        mime_type: None,
1532                        size: None,
1533                        title: None,
1534                        meta: None,
1535                    }),
1536                    " mean?".into(),
1537                ],
1538                cx,
1539            )
1540        });
1541        let send = cx.foreground_executor().spawn(send);
1542        cx.run_until_parked();
1543
1544        model.send_last_completion_stream_text_chunk("Lorem.");
1545        model.end_last_completion_stream();
1546        cx.run_until_parked();
1547        summary_model
1548            .send_last_completion_stream_text_chunk(&format!("Explaining {}", path!("/a/b.md")));
1549        summary_model.end_last_completion_stream();
1550
1551        send.await.unwrap();
1552        let uri = MentionUri::File {
1553            abs_path: path!("/a/b.md").into(),
1554        }
1555        .to_uri();
1556        acp_thread.read_with(cx, |thread, cx| {
1557            assert_eq!(
1558                thread.to_markdown(cx),
1559                formatdoc! {"
1560                    ## User
1561
1562                    What does [@b.md]({uri}) mean?
1563
1564                    ## Assistant
1565
1566                    Lorem.
1567
1568                "}
1569            )
1570        });
1571
1572        cx.run_until_parked();
1573
1574        // Drop the ACP thread, which should cause the session to be dropped as well.
1575        cx.update(|_| {
1576            drop(thread);
1577            drop(acp_thread);
1578        });
1579        agent.read_with(cx, |agent, _| {
1580            assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
1581        });
1582
1583        // Ensure the thread can be reloaded from disk.
1584        assert_eq!(
1585            history_entries(&history_store, cx),
1586            vec![(
1587                HistoryEntryId::AcpThread(session_id.clone()),
1588                format!("Explaining {}", path!("/a/b.md"))
1589            )]
1590        );
1591        let acp_thread = agent
1592            .update(cx, |agent, cx| agent.open_thread(session_id.clone(), cx))
1593            .await
1594            .unwrap();
1595        acp_thread.read_with(cx, |thread, cx| {
1596            assert_eq!(
1597                thread.to_markdown(cx),
1598                formatdoc! {"
1599                    ## User
1600
1601                    What does [@b.md]({uri}) mean?
1602
1603                    ## Assistant
1604
1605                    Lorem.
1606
1607                "}
1608            )
1609        });
1610    }
1611
1612    fn history_entries(
1613        history: &Entity<HistoryStore>,
1614        cx: &mut TestAppContext,
1615    ) -> Vec<(HistoryEntryId, String)> {
1616        history.read_with(cx, |history, _| {
1617            history
1618                .entries()
1619                .map(|e| (e.id(), e.title().to_string()))
1620                .collect::<Vec<_>>()
1621        })
1622    }
1623
1624    fn init_test(cx: &mut TestAppContext) {
1625        env_logger::try_init().ok();
1626        cx.update(|cx| {
1627            let settings_store = SettingsStore::test(cx);
1628            cx.set_global(settings_store);
1629
1630            LanguageModelRegistry::test(cx);
1631        });
1632    }
1633}